From b4c0779e6b971ef6e68a014820efc68753a0544e Mon Sep 17 00:00:00 2001
From: TomDev5 <addr.toml@gmail.com>
Date: Fri, 14 Jun 2024 16:24:56 +0300
Subject: [PATCH] tcp timestamp repr functions

---
 src/socket/tcp.rs | 31 +++++++++++--------------------
 src/wire/mod.rs   |  2 +-
 src/wire/tcp.rs   | 21 +++++++++++++--------
 3 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/src/socket/tcp.rs b/src/socket/tcp.rs
index 4b92125f3..ff45a2950 100644
--- a/src/socket/tcp.rs
+++ b/src/socket/tcp.rs
@@ -14,7 +14,7 @@ use crate::storage::{Assembler, RingBuffer};
 use crate::time::{Duration, Instant};
 use crate::wire::{
     IpAddress, IpEndpoint, IpListenEndpoint, IpProtocol, IpRepr, TcpControl, TcpRepr, TcpSeqNumber,
-    TcpTimestampRepr, TCP_HEADER_LEN,
+    TcpTimestampGenerator, TcpTimestampRepr, TCP_HEADER_LEN,
 };
 
 mod congestion;
@@ -483,7 +483,7 @@ pub struct Socket<'a> {
     congestion_controller: congestion::AnyController,
 
     /// tsval generator - if some, tcp timestamp is enabled
-    tsval_generator: Option<fn() -> u32>,
+    tsval_generator: Option<TcpTimestampGenerator>,
 
     /// 0 if not seen or timestamp not enabled
     last_remote_tsval: u32,
@@ -558,7 +558,7 @@ impl<'a> Socket<'a> {
     }
 
     /// Enable or disable TCP Timestamp.
-    pub fn set_tsval_generator(&mut self, generator: Option<fn() -> u32>) {
+    pub fn set_tsval_generator(&mut self, generator: Option<TcpTimestampGenerator>) {
         self.tsval_generator = generator;
     }
 
@@ -1349,14 +1349,9 @@ impl<'a> Socket<'a> {
 
     fn ack_reply(&mut self, ip_repr: &IpRepr, repr: &TcpRepr) -> (IpRepr, TcpRepr<'static>) {
         let (mut ip_reply_repr, mut reply_repr) = Self::reply(ip_repr, repr);
-
-        reply_repr.timestamp = match (repr.timestamp, self.tsval_generator) {
-            (Some(other_timestamp), Some(tsval_generator)) => Some(TcpTimestampRepr::new(
-                tsval_generator(),
-                other_timestamp.get_tsval(),
-            )),
-            _ => None,
-        };
+        reply_repr.timestamp = repr
+            .timestamp
+            .and_then(|tcp_ts| tcp_ts.generate_reply(self.tsval_generator));
 
         // From RFC 793:
         // [...] an empty acknowledgment segment containing the current send-sequence number
@@ -1991,7 +1986,7 @@ impl<'a> Socket<'a> {
 
         // update last remote tsval
         if let Some(timestamp) = repr.timestamp {
-            self.last_remote_tsval = timestamp.get_tsval();
+            self.last_remote_tsval = timestamp.tsval;
         }
 
         let payload_len = payload.len();
@@ -2286,14 +2281,10 @@ impl<'a> Socket<'a> {
             max_seg_size: None,
             sack_permitted: false,
             sack_ranges: [None, None, None],
-            timestamp: if let Some(ts_generator) = self.tsval_generator {
-                Some(TcpTimestampRepr::new(
-                    ts_generator(),
-                    self.last_remote_tsval,
-                ))
-            } else {
-                None
-            },
+            timestamp: TcpTimestampRepr::generate_reply_with_tsval(
+                self.tsval_generator,
+                self.last_remote_tsval,
+            ),
             payload: &[],
         };
 
diff --git a/src/wire/mod.rs b/src/wire/mod.rs
index 60a67cb49..a552b63b4 100644
--- a/src/wire/mod.rs
+++ b/src/wire/mod.rs
@@ -265,7 +265,7 @@ pub use self::udp::{Packet as UdpPacket, Repr as UdpRepr, HEADER_LEN as UDP_HEAD
 
 pub use self::tcp::{
     Control as TcpControl, Packet as TcpPacket, Repr as TcpRepr, SeqNumber as TcpSeqNumber,
-    TcpOption, TcpTimestampRepr, HEADER_LEN as TCP_HEADER_LEN,
+    TcpOption, TcpTimestampGenerator, TcpTimestampRepr, HEADER_LEN as TCP_HEADER_LEN,
 };
 
 #[cfg(feature = "proto-dhcpv4")]
diff --git a/src/wire/tcp.rs b/src/wire/tcp.rs
index 048eafab6..027e1dad9 100644
--- a/src/wire/tcp.rs
+++ b/src/wire/tcp.rs
@@ -823,10 +823,12 @@ pub struct Repr<'a> {
     pub payload: &'a [u8],
 }
 
+pub type TcpTimestampGenerator = fn() -> u32;
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct TcpTimestampRepr {
-    tsval: u32,
-    tsecr: u32,
+    pub tsval: u32,
+    pub tsecr: u32,
 }
 
 impl TcpTimestampRepr {
@@ -834,12 +836,15 @@ impl TcpTimestampRepr {
         Self { tsval, tsecr }
     }
 
-    pub fn get_tsval(&self) -> u32 {
-        self.tsval
+    pub fn generate_reply(&self, generator: Option<TcpTimestampGenerator>) -> Option<Self> {
+        Self::generate_reply_with_tsval(generator, self.tsval)
     }
 
-    pub fn get_tsecr(&self) -> u32 {
-        self.tsecr
+    pub fn generate_reply_with_tsval(
+        generator: Option<TcpTimestampGenerator>,
+        tsval: u32,
+    ) -> Option<Self> {
+        Some(Self::new(generator?(), tsval))
     }
 }
 
@@ -1023,8 +1028,8 @@ impl<'a> Repr<'a> {
             if let Some(timestamp) = self.timestamp {
                 let tmp = options;
                 options = TcpOption::TimeStamp {
-                    tsval: timestamp.get_tsval(),
-                    tsecr: timestamp.get_tsecr(),
+                    tsval: timestamp.tsval,
+                    tsecr: timestamp.tsecr,
                 }
                 .emit(tmp);
             }