From 197d53c0ab1edbf8545d2ab827a930d66a1ab528 Mon Sep 17 00:00:00 2001
From: Ruihan Li <lrh2000@pku.edu.cn>
Date: Fri, 15 Nov 2024 11:12:34 +0800
Subject: [PATCH] Accept sockets in the ESTABLISHED state

---
 kernel/src/net/socket/ip/stream/connected.rs | 22 -------------------
 kernel/src/net/socket/ip/stream/listen.rs    | 23 +++++++++++++++-----
 2 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/kernel/src/net/socket/ip/stream/connected.rs b/kernel/src/net/socket/ip/stream/connected.rs
index ca5de834f..07e8ceca2 100644
--- a/kernel/src/net/socket/ip/stream/connected.rs
+++ b/kernel/src/net/socket/ip/stream/connected.rs
@@ -89,9 +89,6 @@ impl ConnectedStream {
             Ok(Err(e)) => Err(e),
             Err(RecvError::Finished) => Ok(0),
             Err(RecvError::InvalidState) => {
-                if self.before_established() {
-                    return_errno_with_message!(Errno::EAGAIN, "the connection is not established");
-                }
                 return_errno_with_message!(Errno::ECONNRESET, "the connection is reset")
             }
         }
@@ -110,9 +107,6 @@ impl ConnectedStream {
             Ok(Ok(sent_bytes)) => Ok(sent_bytes),
             Ok(Err(e)) => Err(e),
             Err(SendError::InvalidState) => {
-                if self.before_established() {
-                    return_errno_with_message!(Errno::EAGAIN, "the connection is not established");
-                }
                 // FIXME: `EPIPE` is another possibility, which means that the socket is shut down
                 // for writing. In that case, we should also trigger a `SIGPIPE` if `MSG_NOSIGNAL`
                 // is not specified.
@@ -185,22 +179,6 @@ impl ConnectedStream {
     pub(super) fn set_observer(&self, observer: Weak<dyn SocketEventObserver>) {
         self.bound_socket.set_observer(observer)
     }
-
-    /// Returns whether the connection is before established.
-    ///
-    /// Note that a newly accepted socket may not yet be in the [`TcpState::Established`] state.
-    /// The accept syscall only verifies that a connection request is incoming by ensuring
-    /// that the backlog socket is not in the [`TcpState::Listen`] state.
-    /// However, the socket might still be waiting for further ACKs to complete the establishment process.
-    /// Therefore, it could be in either the [`TcpState::SynSent`] or [`TcpState::SynReceived`] states.
-    /// We must wait until the socket reaches the established state before it can send and receive data.
-    ///
-    /// FIXME: Should we check established state in accept or here?
-    fn before_established(&self) -> bool {
-        self.bound_socket.raw_with(|socket| {
-            socket.state() == TcpState::SynSent || socket.state() == TcpState::SynReceived
-        })
-    }
 }
 
 /// Checks if the peer socket has closed its sending side.
diff --git a/kernel/src/net/socket/ip/stream/listen.rs b/kernel/src/net/socket/ip/stream/listen.rs
index f701ddff5..8f2e91ca4 100644
--- a/kernel/src/net/socket/ip/stream/listen.rs
+++ b/kernel/src/net/socket/ip/stream/listen.rs
@@ -58,7 +58,7 @@ impl ListenStream {
 
         let index = backlog_sockets
             .iter()
-            .position(|backlog_socket| backlog_socket.is_active())
+            .position(|backlog_socket| backlog_socket.can_accept())
             .ok_or_else(|| {
                 Error::with_message(Errno::EAGAIN, "no pending connection is available")
             })?;
@@ -86,10 +86,12 @@ impl ListenStream {
     }
 
     pub(super) fn update_io_events(&self, pollee: &Pollee) {
-        // The lock should be held to avoid data races
         let backlog_sockets = self.backlog_sockets.read();
 
-        let can_accept = backlog_sockets.iter().any(|socket| socket.is_active());
+        let can_accept = backlog_sockets.iter().any(|socket| socket.can_accept());
+
+        // FIXME: If network packets come in simultaneously, the socket state may change in the
+        // middle. This can cause the wrong I/O events to be added or deleted.
         if can_accept {
             pollee.add_events(IoEvents::IN);
         } else {
@@ -131,8 +133,19 @@ impl BacklogSocket {
         }
     }
 
-    fn is_active(&self) -> bool {
-        self.bound_socket.raw_with(|socket| socket.is_active())
+    /// Returns whether the backlog socket can be `accept`ed.
+    ///
+    /// According to the Linux implementation, assuming the TCP Fast Open mechanism is off, a
+    /// backlog socket becomes ready to be returned in the `accept` system call when the 3-way
+    /// handshake is complete (i.e., when it enters the ESTABLISHED state).
+    ///
+    /// The Linux kernel implementation can be found at
+    /// <https://elixir.bootlin.com/linux/v6.11.8/source/net/ipv4/tcp_input.c#L7304>.
+    //
+    // FIMXE: Some sockets may be dead (e.g., RSTed), and such sockets can never become alive
+    // again. We need to remove them from the backlog sockets.
+    fn can_accept(&self) -> bool {
+        self.bound_socket.raw_with(|socket| socket.may_send())
     }
 
     fn remote_endpoint(&self) -> Option<IpEndpoint> {