reset_timeout for tcp

narrowlink · Mar 3, 2024 · b19a1fa · b19a1fa
1 parent 3e9ea28
commit b19a1fa
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/stream/tcb.rs b/src/stream/tcb.rs
@@ -172,9 +172,6 @@ impl Tcb {
         }
     }
     pub(super) fn change_last_ack(&mut self, ack: u32) {
-        self.timeout
-            .as_mut()
-            .reset(tokio::time::Instant::now() + self.tcp_timeout);
         let distance = ack.wrapping_sub(self.last_ack);
 
         if matches!(self.state, TcpState::Established) {
@@ -194,6 +191,11 @@ impl Tcb {
     pub fn is_send_buffer_full(&self) -> bool {
         self.seq.wrapping_sub(self.last_ack) >= MAX_UNACK
     }
+
+    pub(crate) fn reset_timeout(&mut self) {
+        let deadline = tokio::time::Instant::now() + self.tcp_timeout;
+        self.timeout.as_mut().reset(deadline);
+    }
 }
 
 #[derive(Debug, Clone)]

diff --git a/src/stream/tcp.rs b/src/stream/tcp.rs
@@ -209,6 +209,8 @@ impl AsyncRead for IpStackTcpStream {
                 return std::task::Poll::Ready(Err(Error::from(ErrorKind::TimedOut)));
             }
 
+            self.tcb.reset_timeout();
+
             if matches!(self.tcb.get_state(), TcpState::SynReceived(false)) {
                 let flags = tcp_flags::SYN | tcp_flags::ACK;
                 self.packet_to_send = Some(self.create_rev_packet(flags, TTL, None, Vec::new())?);
@@ -412,6 +414,8 @@ impl AsyncWrite for IpStackTcpStream {
         cx: &mut std::task::Context<'_>,
         buf: &[u8],
     ) -> std::task::Poll<Result<usize, std::io::Error>> {
+        self.tcb.reset_timeout();
+
         if (self.tcb.send_window as u64) < self.tcb.avg_send_window.0 / 2
             || self.tcb.is_send_buffer_full()
         {