tcp/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
//! A TCP implementation with a somewhat-BSD-like socket API. It is written as a
//! ["sans-I/O"][sans-io] library meaning it doesn't do any networking I/O itself, it just accepts
//! bytes and provides bytes. A [dependencies](Dependencies) object must be be provided to support
//! setting timers and getting the current time. The TCP state object should probably be used with a
//! reference-counting wrapper so that a reference to the state object can be stored in the timer
//! callbacks.
//!
//! [sans-io]: https://sans-io.readthedocs.io
//!
//! ```
//! use std::cell::RefCell;
//! use std::rc::{Rc, Weak};
//!
//! #[derive(Debug)]
//! struct TcpDependencies {
//!     // a reference to the tcp state
//!     state: Weak<RefCell<tcp::TcpState<Self>>>,
//! }
//!
//! impl tcp::Dependencies for TcpDependencies {
//!     type Instant = std::time::Instant;
//!     type Duration = std::time::Duration;
//!
//!     fn register_timer(
//!         &self,
//!         time: Self::Instant,
//!         f: impl FnOnce(&mut tcp::TcpState<Self>, tcp::TimerRegisteredBy) + Send + Sync + 'static,
//!     ) {
//!         let tcp_state = self.state.upgrade().unwrap();
//!
//!         // TODO: To register timers, you would likely want to involve an async
//!         // runtime. A simple example would create a new task for each timer. The
//!         // task would sleep for some duration and then run the callback.
//!     }
//!
//!     fn current_time(&self) -> Self::Instant {
//!         std::time::Instant::now()
//!     }
//!
//!     fn fork(&self) -> Self {
//!         // TODO: the implementation here would depend on the implementation
//!         // of `register_timer`
//!         todo!();
//!     }
//! }
//!
//! // create the TCP state object
//! let tcp_state = Rc::new_cyclic(|weak| {
//!     let dependencies = TcpDependencies {
//!         state: weak.clone(),
//!     };
//!     RefCell::new(tcp::TcpState::new(dependencies, tcp::TcpConfig::default()))
//! });
//!
//! let mut tcp_state = tcp_state.borrow_mut();
//!
//! // connect to port 80
//! let dst_addr = "127.0.0.1:80".parse().unwrap();
//! tcp_state.connect(dst_addr, || {
//!     // here we would ask the network interface for an unused port (implicit bind),
//!     // or where we would use the port assigned to a raw IP socket
//!     let bind_addr = "127.0.0.1:2532".parse().unwrap();
//!     Ok::<_, ()>((bind_addr, ()))
//! }).unwrap();
//!
//! // get the SYN packet from the connect
//! let (header, _payload) = tcp_state.pop_packet().unwrap();
//! assert!(header.flags.contains(tcp::TcpFlags::SYN));
//! assert_eq!(header.dst(), dst_addr);
//! ```

// There are three related state types in this crate:
//
// - `TcpState` — This is the public-facing type for the TCP state. Its methods take shared or
//   mutable references. It contains a non-public `TcpStateEnum`.
// - `TcpStateEnum` — An enum of all individual TCP state types (ex: `ListeningState`,
//   `EstablishedState`). It implements the `TcpStateTrait` trait, so its methods usually take owned
//   objects and return owned objects.
// - `TcpStateTrait` — A trait implemented by each individual TCP state type, as well as the
//   `TcpStateEnum` enum that encapsulates all individual states. Its methods usually take owned
//   state objects and return owned `TcpStateEnum` objects.

#![forbid(unsafe_code)]

use std::fmt::Debug;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, SocketAddrV4};

use bytes::{Bytes, BytesMut};

pub mod util;

mod buffer;
mod connection;
mod seq;
mod states;
mod window_scaling;

#[cfg(test)]
mod tests;

use crate::states::{
    CloseWaitState, ClosedState, ClosingState, EstablishedState, FinWaitOneState, FinWaitTwoState,
    InitState, LastAckState, ListenState, RstState, SynReceivedState, SynSentState, TimeWaitState,
};
use crate::util::SmallArrayBackedSlice;

/// A collection of methods that allow the TCP state to interact with the external system.
pub trait Dependencies: Debug + Sized {
    type Instant: crate::util::time::Instant<Duration = Self::Duration>;
    type Duration: crate::util::time::Duration;

    /// Register a timer. The callback will be run on the parent [state](TcpState). The callback can
    /// use the [`TimerRegisteredBy`] argument to know whether the timer was registered by the
    /// parent state or one of its child states.
    ///
    /// If a child state has not yet been accept()ed, it will be owned by a parent state. When a
    /// child state registers a timer, the timer's callback will run on the parent state and the
    /// callback will be given the `TimerRegisteredBy::Child` argument so that the callback can
    /// delegate accordingly.
    fn register_timer(
        &self,
        time: Self::Instant,
        f: impl FnOnce(&mut TcpState<Self>, TimerRegisteredBy) + Send + Sync + 'static,
    );

    /// Get the current time.
    fn current_time(&self) -> Self::Instant;

    /// Create a new `Dependencies` for use by a child state. When a timer is registered by the
    /// child state using this new object, the timer's callback will be run on the parent's state
    /// with the `TimerRegisteredBy::Child` argument so that the parent knows to run the callback on
    /// one of its child states.
    ///
    /// When a child state has been accept()ed, it will no longer be owned by the parent state and
    /// the parent state has no way to access this child state. The child state's `Dependencies`
    /// should be updated during the [`finalize`](AcceptedTcpState::finalize) call (on the
    /// [`AcceptedTcpState`] returned from [`accept`](TcpState::accept)) to run callbacks directly
    /// on this state instead, and the callbacks should be given `TimerRegisteredBy::Parent` (the
    /// child state has effectively become a parent state). This `Dependencies` object should also
    /// make sure that all existing timers from before the state was accept()ed are also updated to
    /// run callbacks directly on the state.
    fn fork(&self) -> Self;
}

/// Specifies whether the callback is meant to run on the parent state or a child state.
///
/// For example if a child state registers a timer, a value of `TimerRegisteredBy::Child` will be
/// given to the callback so that it knows to apply the callback to a child state, not the parent
/// state.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum TimerRegisteredBy {
    Parent,
    Child,
}

#[enum_dispatch::enum_dispatch]
trait TcpStateTrait<X>: Debug + Sized
where
    X: Dependencies,
    TcpStateEnum<X>: From<Self>,
{
    /// Start closing this socket. It may or may not close immediately depending on what state the
    /// socket is currently in.
    fn close(self) -> (TcpStateEnum<X>, Result<(), CloseError>) {
        (self.into(), Err(CloseError::InvalidState))
    }

    /// Start closing this socket by sending an RST packet. It may or may not close immediately
    /// depending on what state the socket is currently in.
    ///
    /// TODO:
    /// RFC 9293: "The side of a connection issuing a reset should enter the TIME-WAIT state, [...]"
    fn rst_close(self) -> (TcpStateEnum<X>, Result<(), RstCloseError>) {
        (self.into(), Err(RstCloseError::InvalidState))
    }

    fn shutdown(self, _how: Shutdown) -> (TcpStateEnum<X>, Result<(), ShutdownError>) {
        (self.into(), Err(ShutdownError::InvalidState))
    }

    fn listen<T, E>(
        self,
        _backlog: u32,
        _associate_fn: impl FnOnce() -> Result<T, E>,
    ) -> (TcpStateEnum<X>, Result<T, ListenError<E>>) {
        (self.into(), Err(ListenError::InvalidState))
    }

    fn connect<T, E>(
        self,
        _addr: SocketAddrV4,
        _associate_fn: impl FnOnce() -> Result<(SocketAddrV4, T), E>,
    ) -> (TcpStateEnum<X>, Result<T, ConnectError<E>>) {
        (self.into(), Err(ConnectError::InvalidState))
    }

    /// Accept a new child state from the pending connection queue. The TCP state for the child
    /// socket is returned. The [`AcceptedTcpState::finalize`] method must be called immediately on
    /// the returned child state before any code calls into the parent state again, otherwise the
    /// child may miss some timer events.
    fn accept(self) -> (TcpStateEnum<X>, Result<AcceptedTcpState<X>, AcceptError>) {
        (self.into(), Err(AcceptError::InvalidState))
    }

    fn send(self, _reader: impl Read, _len: usize) -> (TcpStateEnum<X>, Result<usize, SendError>) {
        (self.into(), Err(SendError::InvalidState))
    }

    fn recv(self, _writer: impl Write, _len: usize) -> (TcpStateEnum<X>, Result<usize, RecvError>) {
        (self.into(), Err(RecvError::InvalidState))
    }

    /// Returns the number of bytes added to the TCP state's receive buffer. This may be
    /// smaller (ex: duplicate packet) or larger (ex: there is a non-empty reassembly queue)
    /// than the packet payload length.
    fn push_packet(
        self,
        _header: &TcpHeader,
        _payload: Payload,
    ) -> (TcpStateEnum<X>, Result<u32, PushPacketError>) {
        (self.into(), Err(PushPacketError::InvalidState))
    }

    fn pop_packet(
        self,
    ) -> (
        TcpStateEnum<X>,
        Result<(TcpHeader, Payload), PopPacketError>,
    ) {
        (self.into(), Err(PopPacketError::InvalidState))
    }

    fn clear_error(&mut self) -> Option<TcpError>;

    fn poll(&self) -> PollState;

    fn wants_to_send(&self) -> bool;

    fn local_remote_addrs(&self) -> Option<(SocketAddrV4, SocketAddrV4)>;
}

#[derive(Debug)]
pub struct TcpState<X: Dependencies>(Option<TcpStateEnum<X>>);

// this exposes many of the methods from `TcpStateTrait`, but not necessarily all of them (for
// example we don't expose `rst_close()`).
impl<X: Dependencies> TcpState<X> {
    pub fn new(deps: X, config: TcpConfig) -> Self {
        let new_state = InitState::new(deps, config);
        Self(Some(new_state.into()))
    }

    #[inline]
    fn with_state<T>(&mut self, f: impl FnOnce(TcpStateEnum<X>) -> (TcpStateEnum<X>, T)) -> T {
        // get the current state, pass it to `f`, and then put it back (`f` may actually replace the
        // state with an entirely different state object)
        let state = self.0.take().unwrap();
        let (state, rv) = f(state);
        self.0 = Some(state);

        rv
    }

    #[inline]
    pub fn close(&mut self) -> Result<(), CloseError> {
        self.with_state(|state| state.close())
    }

    #[inline]
    pub fn shutdown(&mut self, how: Shutdown) -> Result<(), ShutdownError> {
        self.with_state(|state| state.shutdown(how))
    }

    #[inline]
    pub fn listen<T, E>(
        &mut self,
        backlog: u32,
        associate_fn: impl FnOnce() -> Result<T, E>,
    ) -> Result<T, ListenError<E>> {
        self.with_state(|state| state.listen(backlog, associate_fn))
    }

    #[inline]
    pub fn connect<T, E>(
        &mut self,
        addr: SocketAddrV4,
        associate_fn: impl FnOnce() -> Result<(SocketAddrV4, T), E>,
    ) -> Result<T, ConnectError<E>> {
        self.with_state(|state| state.connect(addr, associate_fn))
    }

    #[inline]
    pub fn accept(&mut self) -> Result<AcceptedTcpState<X>, AcceptError> {
        self.with_state(|state| state.accept())
    }

    #[inline]
    pub fn send(&mut self, reader: impl Read, len: usize) -> Result<usize, SendError> {
        self.with_state(|state| state.send(reader, len))
    }

    #[inline]
    pub fn recv(&mut self, writer: impl Write, len: usize) -> Result<usize, RecvError> {
        self.with_state(|state| state.recv(writer, len))
    }

    #[inline]
    pub fn push_packet(
        &mut self,
        header: &TcpHeader,
        payload: Payload,
    ) -> Result<u32, PushPacketError> {
        self.with_state(|state| state.push_packet(header, payload))
    }

    #[inline]
    pub fn pop_packet(&mut self) -> Result<(TcpHeader, Payload), PopPacketError> {
        self.with_state(|state| state.pop_packet())
    }

    #[inline]
    pub fn clear_error(&mut self) -> Option<TcpError> {
        self.0.as_mut().unwrap().clear_error()
    }

    #[inline]
    pub fn poll(&self) -> PollState {
        self.0.as_ref().unwrap().poll()
    }

    #[inline]
    pub fn wants_to_send(&self) -> bool {
        self.0.as_ref().unwrap().wants_to_send()
    }

    #[inline]
    pub fn local_remote_addrs(&self) -> Option<(SocketAddrV4, SocketAddrV4)> {
        self.0.as_ref().unwrap().local_remote_addrs()
    }
}

/// A macro that forwards an argument-less method to the inner type.
///
/// ```ignore
/// // forward!(as_init, Option<&InitState<X>>);
/// #[inline]
/// pub fn as_init(&self) -> Option<&InitState<X>> {
///     self.0.as_ref().unwrap().as_init()
/// }
/// ```
#[cfg(test)]
macro_rules! forward {
    ($fn_name:ident, $($return_type:tt)*) => {
        #[inline]
        pub fn $fn_name(&self) -> $($return_type)* {
            self.0.as_ref().unwrap().$fn_name()
        }
    };
}

#[cfg(test)]
impl<X: Dependencies> TcpState<X> {
    forward!(as_init, Option<&InitState<X>>);
    forward!(as_listen, Option<&ListenState<X>>);
    forward!(as_syn_sent, Option<&SynSentState<X>>);
    forward!(as_syn_received, Option<&SynReceivedState<X>>);
    forward!(as_established, Option<&EstablishedState<X>>);
    forward!(as_fin_wait_one, Option<&FinWaitOneState<X>>);
    forward!(as_fin_wait_two, Option<&FinWaitTwoState<X>>);
    forward!(as_closing, Option<&ClosingState<X>>);
    forward!(as_time_wait, Option<&TimeWaitState<X>>);
    forward!(as_close_wait, Option<&CloseWaitState<X>>);
    forward!(as_last_ack, Option<&LastAckState<X>>);
    forward!(as_rst, Option<&RstState<X>>);
    forward!(as_closed, Option<&ClosedState<X>>);
}

#[enum_dispatch::enum_dispatch(TcpStateTrait<X>)]
#[derive(Debug)]
enum TcpStateEnum<X: Dependencies> {
    Init(InitState<X>),
    Listen(ListenState<X>),
    SynSent(SynSentState<X>),
    SynReceived(SynReceivedState<X>),
    Established(EstablishedState<X>),
    FinWaitOne(FinWaitOneState<X>),
    FinWaitTwo(FinWaitTwoState<X>),
    Closing(ClosingState<X>),
    TimeWait(TimeWaitState<X>),
    CloseWait(CloseWaitState<X>),
    LastAck(LastAckState<X>),
    Rst(RstState<X>),
    Closed(ClosedState<X>),
}

/// A macro that creates a method which casts to an inner variant.
///
/// ```ignore
/// // as_impl!(as_init, Init, InitState);
/// #[inline]
/// pub fn as_init(&self) -> Option<&InitState<X>> {
///     match self {
///         Self::Init(x) => Some(x),
///         _ => None,
///     }
/// }
/// ```
#[cfg(test)]
macro_rules! as_impl {
    ($fn_name:ident, $variant:ident, $return_type:ident) => {
        #[inline]
        pub fn $fn_name(&self) -> Option<&$return_type<X>> {
            match self {
                Self::$variant(x) => Some(x),
                _ => None,
            }
        }
    };
}

/// Casts to concrete types. This should only be called from unit tests to verify state.
#[cfg(test)]
impl<X: Dependencies> TcpStateEnum<X> {
    as_impl!(as_init, Init, InitState);
    as_impl!(as_listen, Listen, ListenState);
    as_impl!(as_syn_sent, SynSent, SynSentState);
    as_impl!(as_syn_received, SynReceived, SynReceivedState);
    as_impl!(as_established, Established, EstablishedState);
    as_impl!(as_fin_wait_one, FinWaitOne, FinWaitOneState);
    as_impl!(as_fin_wait_two, FinWaitTwo, FinWaitTwoState);
    as_impl!(as_closing, Closing, ClosingState);
    as_impl!(as_time_wait, TimeWait, TimeWaitState);
    as_impl!(as_close_wait, CloseWait, CloseWaitState);
    as_impl!(as_last_ack, LastAck, LastAckState);
    as_impl!(as_rst, Rst, RstState);
    as_impl!(as_closed, Closed, ClosedState);
}

/// An accept()ed TCP state. This is used to ensure that the caller uses
/// [`finalize`](Self::finalize) to update the state's `Dependencies` since the state is no longer
/// owned by the listening socket.
// we use a wrapper struct around an enum so that public code can't access the inner state object
pub struct AcceptedTcpState<X: Dependencies>(AcceptedTcpStateInner<X>);

/// An "established" or "close-wait" TCP state can be accept()ed, so we need to be able to store
/// either state.
enum AcceptedTcpStateInner<X: Dependencies> {
    Established(EstablishedState<X>),
    CloseWait(CloseWaitState<X>),
}

impl<X: Dependencies> AcceptedTcpState<X> {
    /// This allows the caller to update the state's `Dependencies`.
    ///
    /// This must be called immediately after [`TcpState::accept`], otherwise the accept()ed socket
    /// may miss some of its timer events.
    pub fn finalize(mut self, f: impl FnOnce(&mut X)) -> TcpState<X> {
        let deps = match &mut self.0 {
            AcceptedTcpStateInner::Established(state) => &mut state.common.deps,
            AcceptedTcpStateInner::CloseWait(state) => &mut state.common.deps,
        };

        // allow the caller to update `deps` for this state since this state is changing owners
        f(deps);

        TcpState(Some(self.0.into()))
    }

    pub fn local_addr(&self) -> SocketAddrV4 {
        match &self.0 {
            AcceptedTcpStateInner::Established(state) => state.connection.local_addr,
            AcceptedTcpStateInner::CloseWait(state) => state.connection.local_addr,
        }
    }

    pub fn remote_addr(&self) -> SocketAddrV4 {
        match &self.0 {
            AcceptedTcpStateInner::Established(state) => state.connection.remote_addr,
            AcceptedTcpStateInner::CloseWait(state) => state.connection.remote_addr,
        }
    }
}

impl<X: Dependencies> TryFrom<TcpStateEnum<X>> for AcceptedTcpState<X> {
    type Error = TcpStateEnum<X>;

    fn try_from(state: TcpStateEnum<X>) -> Result<Self, Self::Error> {
        match state {
            TcpStateEnum::Established(state) => Ok(Self(AcceptedTcpStateInner::Established(state))),
            TcpStateEnum::CloseWait(state) => Ok(Self(AcceptedTcpStateInner::CloseWait(state))),
            // return the state back to the caller
            state => Err(state),
        }
    }
}

impl<X: Dependencies> From<AcceptedTcpStateInner<X>> for TcpStateEnum<X> {
    fn from(inner: AcceptedTcpStateInner<X>) -> Self {
        match inner {
            AcceptedTcpStateInner::Established(state) => state.into(),
            AcceptedTcpStateInner::CloseWait(state) => state.into(),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Shutdown {
    Read,
    Write,
    Both,
}

#[derive(Debug)]
pub enum TcpError {
    ResetSent,
    ResetReceived,
    /// The connection was closed while it was connecting, and no RST was sent or received.
    ClosedWhileConnecting,
    TimedOut,
}

// errors for operations on `TcpStateTrait` objects

#[derive(Debug)]
pub enum CloseError {
    InvalidState,
}

#[derive(Debug)]
enum RstCloseError {
    InvalidState,
}

#[derive(Debug)]
pub enum ListenError<E> {
    InvalidState,
    FailedAssociation(E),
}

#[derive(Debug)]
pub enum ConnectError<E> {
    InvalidState,
    /// A previous connection attempt is in progress.
    InProgress,
    /// A connection has previously been attempted and was either successful or unsuccessful (it may
    /// or may not have reached the "established" state). The connection may be established, timed
    /// out, closing, half-closed, closed, etc. This does not include connection attempts that are
    /// in progress ("syn-sent" or "syn-received" states).
    AlreadyConnected,
    /// Is already listening for new connections.
    IsListening,
    FailedAssociation(E),
}

#[derive(Debug)]
pub enum AcceptError {
    InvalidState,
    NothingToAccept,
}

#[derive(Debug)]
pub enum ShutdownError {
    NotConnected,
    InvalidState,
}

#[derive(Debug)]
pub enum SendError {
    InvalidState,
    Full,
    NotConnected,
    StreamClosed,
    Io(std::io::Error),
}

#[derive(Debug)]
pub enum RecvError {
    InvalidState,
    Empty,
    NotConnected,
    /// The peer has sent a FIN, so no more data will be received.
    StreamClosed,
    Io(std::io::Error),
}

#[derive(Debug)]
pub enum PushPacketError {
    InvalidState,
}

#[derive(Debug)]
pub enum PopPacketError {
    InvalidState,
    NoPacket,
}

// segment/packet headers

bitflags::bitflags! {
    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
    pub struct PollState: u32 {
        /// Data can be read.
        const READABLE = 1 << 0;
        /// Data can be written.
        const WRITABLE = 1 << 1;
        /// There is a pending error that can be read using [`TcpState::clear_error`].
        const ERROR = 1 << 2;
        /// The connection has been closed for receiving. This is not mutually exclusive with
        /// `READABLE` (even if it's closed for receiving, there may still be buffered data to
        /// read). Some possible causes are:
        /// - Received a FIN packet.
        /// - Sent or received a RST packet.
        /// - TCP was closed.
        const RECV_CLOSED = 1 << 3;
        /// The connection has been closed for sending. This should be mutually exclusive with
        /// `WRITABLE` (there would be no point in writing data if it's closed for sending). Some
        /// possible causes are:
        /// - Sent a FIN packet.
        /// - Sent or received a RST packet.
        /// - TCP was `shutdown()` for writing.
        /// - TCP was closed.
        const SEND_CLOSED = 1 << 4;
        /// Is listening for new connections.
        const LISTENING = 1 << 5;
        /// A listening socket has a new incoming connection that can be accepted.
        const READY_TO_ACCEPT = 1 << 6;
        /// Connection is in the process of opening. More specifically this means that it is in
        /// either the "syn-sent" or "syn-received" states.
        const CONNECTING = 1 << 7;
        /// A connection has previously been attempted and was either successful or unsuccessful (it
        /// may or may not have reached the "established" state). The connection may be established,
        /// timed out, closing, half-closed, closed, etc. This does not include connection attempts
        /// that are in progress ("syn-sent" or "syn-received" states).
        const CONNECTED = 1 << 8;
        /// TCP is fully closed (in the "closed" state). This may not be set immediately after a
        /// `close()` call, for example if `close()` was called while in the "established" state,
        /// and now is in the "fin-wait-1" state. This does not include the initial state (we don't
        /// consider a new TCP to be "closed").
        const CLOSED = 1 << 9;
    }
}

#[derive(Copy, Clone, Debug)]
pub struct TcpConfig {
    pub(crate) window_scaling_enabled: bool,
}

impl TcpConfig {
    pub fn window_scaling(&mut self, enable: bool) {
        self.window_scaling_enabled = enable;
    }
}

impl Default for TcpConfig {
    fn default() -> Self {
        Self {
            window_scaling_enabled: true,
        }
    }
}

bitflags::bitflags! {
    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
    pub struct TcpFlags: u8 {
        const FIN = 1 << 0;
        const SYN = 1 << 1;
        const RST = 1 << 2;
        const PSH = 1 << 3;
        const ACK = 1 << 4;
        const URG = 1 << 5;
        const ECE = 1 << 6;
        const CWR = 1 << 7;
    }
}

#[derive(Copy, Clone, Debug)]
pub struct TcpHeader {
    pub ip: Ipv4Header,
    pub flags: TcpFlags,
    pub src_port: u16,
    pub dst_port: u16,
    pub seq: u32,
    pub ack: u32,
    pub window_size: u16,
    pub selective_acks: Option<SmallArrayBackedSlice<4, (u32, u32)>>,
    pub window_scale: Option<u8>,
    pub timestamp: Option<u32>,
    pub timestamp_echo: Option<u32>,
}

impl TcpHeader {
    pub fn src(&self) -> SocketAddrV4 {
        SocketAddrV4::new(self.ip.src, self.src_port)
    }

    pub fn dst(&self) -> SocketAddrV4 {
        SocketAddrV4::new(self.ip.dst, self.dst_port)
    }
}

#[derive(Copy, Clone, Debug)]
pub struct Ipv4Header {
    pub src: Ipv4Addr,
    pub dst: Ipv4Addr,
}

/// A packet payload containing a list of [byte](Bytes) chunks.
///
/// The sum of the lengths of each chunk must be at most [`u32::MAX`], otherwise operations on the
/// payload or other code using the payload may panic.
// TODO: Intuitively this seems like a good place to use a `SmallVec` to optimize the common case
// where there are a small number of chunks per packet. But I'm leaving this until we can test `Vec`
// vs `SmallVec` in a benchmark to see if there's any performance improvement in practice.
#[derive(Clone, Debug, Default)]
pub struct Payload(pub Vec<Bytes>);

// We don't implement `PartialEq` or `Eq` since it's not clear what equality means. Are payloads
// equal if they just contain the same bytes, or are they equal only if the chunks are exactly the
// same? For example is the payload `["hello", "world"]` the same as `["helloworld"]`?
static_assertions::assert_not_impl_any!(Payload: PartialEq, Eq);

impl Payload {
    /// Returns the number of bytes in the payload.
    pub fn len(&self) -> u32 {
        self.0
            .iter()
            // `fold` rather than `sum` so that we always panic on overflow
            .fold(0usize, |acc, x| acc.checked_add(x.len()).unwrap())
            .try_into()
            .unwrap()
    }

    /// Returns true if the payload has no data (no byte chunks or only empty byte chunks).
    pub fn is_empty(&self) -> bool {
        // should be faster than checking `self.len() == 0`
        self.0.iter().all(|x| x.len() == 0)
    }

    /// Concatenate the byte chunks into a single byte chunk. Unless the payload is empty or has a
    /// single chunk, this will allocate a large buffer and copy all of the individual chunks to
    /// this new buffer.
    pub fn concat(&self) -> Bytes {
        let num_bytes = self.len() as usize;
        let num_chunks = self.0.len();

        if num_bytes == 0 {
            return Bytes::new();
        }

        if num_chunks == 1 {
            // there's only one chunk, so just return a reference to the chunk
            return self.0[0].clone();
        }

        let mut bytes = BytesMut::with_capacity(num_bytes);

        for chunk in &self.0 {
            bytes.extend_from_slice(chunk);
        }

        debug_assert_eq!(bytes.len(), bytes.capacity());

        bytes.freeze()
    }
}

impl From<Bytes> for Payload {
    fn from(bytes: Bytes) -> Self {
        Self(vec![bytes])
    }
}

impl From<BytesMut> for Payload {
    fn from(bytes: BytesMut) -> Self {
        bytes.freeze().into()
    }
}