Simplify the TCP state check

kernel/libs/aster-bigtcp/src/socket/bound/mod.rs

@@ -6,7 +6,7 @@ mod tcp_listen;
 mod udp;
 
 pub use common::NeedIfacePoll;
-pub use tcp_conn::{ConnectState, TcpConnection};
+pub use tcp_conn::{ConnectState, RawTcpSocketExt, TcpConnection};
 pub(crate) use tcp_conn::{TcpConnectionBg, TcpProcessResult};
 pub use tcp_listen::TcpListener;
 pub(crate) use tcp_listen::TcpListenerBg;

kernel/libs/aster-bigtcp/src/socket/bound/tcp_conn.rs

@@ -23,8 +23,7 @@ use crate::{
     socket::{
         event::SocketEvents,
         option::{RawTcpOption, RawTcpSetOption},
-        unbound::new_tcp_socket,
+        unbound::{new_tcp_socket, RawTcpSocket},
-        RawTcpSocket, TcpStateCheck,
     },
     socket_table::ConnectionKey,
 };
@@ -37,7 +36,7 @@ pub struct TcpConnectionInner<E: Ext> {
     connection_key: ConnectionKey,
 }
 
-pub(super) struct RawTcpSocketExt<E: Ext> {
+pub struct RawTcpSocketExt<E: Ext> {
     socket: Box<RawTcpSocket>,
     pub(super) listener: Option<Arc<TcpListenerBg<E>>>,
     has_connected: bool,
@@ -57,6 +56,25 @@ impl<E: Ext> DerefMut for RawTcpSocketExt<E> {
     }
 }
 
+impl<E: Ext> RawTcpSocketExt<E> {
+    /// Checks if the socket may receive any new data.
+    ///
+    /// This is similar to [`RawTcpSocket::may_recv`]. However, this method checks if there can be
+    /// _new_ data. In other words, if there is already buffered data in the socket,
+    /// [`RawTcpSocket::may_recv`] will always return true since it is possible to receive the
+    /// buffered data, but this method may return false if the peer has closed its sending half (so
+    /// no new data can come in).
+    pub fn may_recv_new(&self) -> bool {
+        // See also the implementation of `RawTcpSocket::may_recv`.
+        match self.state() {
+            State::Established => true,
+            // Our sending half is closed, but the peer's sending half is still active.
+            State::FinWait1 | State::FinWait2 => true,
+            _ => false,
+        }
+    }
+}
+
 define_boolean_value!(
     /// Whether the TCP connection became dead.
     TcpConnBecameDead
@@ -67,7 +85,9 @@ impl<E: Ext> RawTcpSocketExt<E> {
         &mut self,
         this: &Arc<TcpConnectionBg<E>>,
     ) -> (SocketEvents, TcpConnBecameDead) {
-        if self.may_send() && !self.has_connected {
+        let may_send = self.may_send();
+
+        if may_send && !self.has_connected {
             self.has_connected = true;
 
             if let Some(ref listener) = self.listener {
@@ -81,13 +101,13 @@ impl<E: Ext> RawTcpSocketExt<E> {
 
         let became_dead = self.check_dead(this);
 
-        let events = if self.is_peer_closed() {
+        let mut events = SocketEvents::empty();
-            SocketEvents::PEER_CLOSED
+        if !self.may_recv_new() {
-        } else if self.is_closed() {
+            events |= SocketEvents::CLOSED_RECV;
-            SocketEvents::CLOSED
+        }
-        } else {
+        if !may_send {
-            SocketEvents::empty()
+            events |= SocketEvents::CLOSED_SEND;
-        };
+        }
 
         (events, became_dead)
     }
@@ -317,7 +337,7 @@ impl<E: Ext> TcpConnection<E> {
 
         socket.listener = None;
 
-        if socket.is_closed() {
+        if socket.state() == State::Closed {
             return false;
         }
 
@@ -333,7 +353,7 @@ impl<E: Ext> TcpConnection<E> {
     // polling time.
     pub fn raw_with<F, R>(&self, f: F) -> R
     where
-        F: FnOnce(&RawTcpSocket) -> R,
+        F: FnOnce(&RawTcpSocketExt<E>) -> R,
     {
         let socket = self.0.inner.lock();
         f(&socket)

kernel/libs/aster-bigtcp/src/socket/bound/tcp_listen.rs

@@ -20,8 +20,7 @@ use crate::{
     iface::{BindPortConfig, BoundPort},
     socket::{
         option::{RawTcpOption, RawTcpSetOption},
-        unbound::new_tcp_socket,
+        unbound::{new_tcp_socket, RawTcpSocket},
-        RawTcpSocket,
     },
     socket_table::{ConnectionKey, ListenerKey},
 };

kernel/libs/aster-bigtcp/src/socket/event.rs

@@ -15,7 +15,9 @@ bitflags::bitflags! {
     pub struct SocketEvents: u8 {
         const CAN_RECV = 1;
         const CAN_SEND = 2;
-        const PEER_CLOSED = 4;
+        /// Receiving new data isn't possible anymore.
-        const CLOSED = 8;
+        const CLOSED_RECV = 4;
+        /// Sending data isn't possible anymore.
+        const CLOSED_SEND = 8;
     }
 }

kernel/libs/aster-bigtcp/src/socket/mod.rs

@@ -3,15 +3,14 @@
 mod bound;
 mod event;
 mod option;
-mod state;
 mod unbound;
 
-pub use bound::{ConnectState, NeedIfacePoll, TcpConnection, TcpListener, UdpSocket};
+pub use bound::{
+    ConnectState, NeedIfacePoll, RawTcpSocketExt, TcpConnection, TcpListener, UdpSocket,
+};
 pub(crate) use bound::{TcpConnectionBg, TcpListenerBg, TcpProcessResult, UdpSocketBg};
 pub use event::{SocketEventObserver, SocketEvents};
 pub use option::{RawTcpOption, RawTcpSetOption};
-pub use state::TcpStateCheck;
+pub use unbound::{
-pub use unbound::{TCP_RECV_BUF_LEN, TCP_SEND_BUF_LEN, UDP_RECV_PAYLOAD_LEN, UDP_SEND_PAYLOAD_LEN};
+    RawUdpSocket, TCP_RECV_BUF_LEN, TCP_SEND_BUF_LEN, UDP_RECV_PAYLOAD_LEN, UDP_SEND_PAYLOAD_LEN,
-
+};
-pub type RawTcpSocket = smoltcp::socket::tcp::Socket<'static>;
-pub type RawUdpSocket = smoltcp::socket::udp::Socket<'static>;

kernel/libs/aster-bigtcp/src/socket/option.rs

@@ -2,7 +2,7 @@
 
 use smoltcp::time::Duration;
 
-use super::{NeedIfacePoll, RawTcpSocket};
+use super::{unbound::RawTcpSocket, NeedIfacePoll};
 
 /// A trait defines setting socket options on a raw socket.
 pub trait RawTcpSetOption {

kernel/libs/aster-bigtcp/src/socket/state.rs

@@ -1,31 +0,0 @@
-// SPDX-License-Identifier: MPL-2.0
-
-pub use smoltcp::socket::tcp::State as TcpState;
-
-use super::RawTcpSocket;
-
-pub trait TcpStateCheck {
-    /// Checks if the peer socket has closed its sending side.
-    ///
-    /// If the sending side of this socket is also closed, this method will return `false`.
-    /// In such cases, you should verify using [`is_closed`].
-    fn is_peer_closed(&self) -> bool;
-
-    /// Checks if the socket is fully closed.
-    ///
-    /// This function returns `true` if both this socket and the peer have closed their sending sides.
-    ///
-    /// This TCP state corresponds to the `Normal Close Sequence` and `Simultaneous Close Sequence`
-    /// as outlined in RFC793 (https://datatracker.ietf.org/doc/html/rfc793#page-39).
-    fn is_closed(&self) -> bool;
-}
-
-impl TcpStateCheck for RawTcpSocket {
-    fn is_peer_closed(&self) -> bool {
-        self.state() == TcpState::CloseWait
-    }
-
-    fn is_closed(&self) -> bool {
-        !self.is_open() || self.state() == TcpState::Closing || self.state() == TcpState::LastAck
-    }
-}

kernel/libs/aster-bigtcp/src/socket/unbound.rs

@@ -2,7 +2,8 @@
 
 use alloc::{boxed::Box, vec};
 
-use super::{RawTcpSocket, RawUdpSocket};
+pub(super) type RawTcpSocket = smoltcp::socket::tcp::Socket<'static>;
+pub type RawUdpSocket = smoltcp::socket::udp::Socket<'static>;
 
 pub(super) fn new_tcp_socket() -> Box<RawTcpSocket> {
     let raw_tcp_socket = {

kernel/src/net/iface/mod.rs

@@ -11,6 +11,8 @@ pub use poll::lazy_init;
 pub type Iface = dyn aster_bigtcp::iface::Iface<ext::BigtcpExt>;
 pub type BoundPort = aster_bigtcp::iface::BoundPort<ext::BigtcpExt>;
 
+pub type RawTcpSocketExt = aster_bigtcp::socket::RawTcpSocketExt<ext::BigtcpExt>;
+
 pub type TcpConnection = aster_bigtcp::socket::TcpConnection<ext::BigtcpExt>;
 pub type TcpListener = aster_bigtcp::socket::TcpListener<ext::BigtcpExt>;
 pub type UdpSocket = aster_bigtcp::socket::UdpSocket<ext::BigtcpExt>;

kernel/src/net/socket/ip/stream/connected.rs

@@ -4,7 +4,7 @@ use core::sync::atomic::{AtomicBool, Ordering};
 
 use aster_bigtcp::{
     errors::tcp::{RecvError, SendError},
-    socket::{NeedIfacePoll, RawTcpSetOption, RawTcpSocket, TcpStateCheck},
+    socket::{NeedIfacePoll, RawTcpSetOption},
     wire::IpEndpoint,
 };
 
@@ -12,7 +12,7 @@ use super::StreamObserver;
 use crate::{
     events::IoEvents,
     net::{
-        iface::{Iface, TcpConnection},
+        iface::{Iface, RawTcpSocketExt, TcpConnection},
         socket::util::{send_recv_flags::SendRecvFlags, shutdown_cmd::SockShutdownCmd},
     },
     prelude::*,
@@ -34,15 +34,8 @@ pub struct ConnectedStream {
     /// connection is established asynchronously will succeed and any subsequent `connect()` will
     /// fail.
     is_new_connection: bool,
-    /// Indicates if the receiving side of this socket is closed.
+    /// Indicates if the receiving side of this socket is shut down by the user.
-    ///
+    is_receiving_shut: AtomicBool,
-    /// The receiving side may be closed if this side disables reading
-    /// or if the peer side closes its sending half.
-    is_receiving_closed: AtomicBool,
-    /// Indicates if the sending side of this socket is closed.
-    ///
-    /// The sending side can only be closed if this side disables writing.
-    is_sending_closed: AtomicBool,
 }
 
 impl ConnectedStream {
@@ -55,8 +48,7 @@ impl ConnectedStream {
             tcp_conn,
             remote_endpoint,
             is_new_connection,
-            is_receiving_closed: AtomicBool::new(false),
+            is_receiving_shut: AtomicBool::new(false),
-            is_sending_closed: AtomicBool::new(false),
         }
     }
 
@@ -64,12 +56,11 @@ impl ConnectedStream {
         let mut events = IoEvents::empty();
 
         if cmd.shut_read() {
-            self.is_receiving_closed.store(true, Ordering::Relaxed);
+            self.is_receiving_shut.store(true, Ordering::Relaxed);
             events |= IoEvents::IN | IoEvents::RDHUP;
         }
 
         if cmd.shut_write() {
-            self.is_sending_closed.store(true, Ordering::Relaxed);
             if !self.tcp_conn.close() {
                 return_errno_with_message!(Errno::ENOTCONN, "the socket is not connected");
             }
@@ -94,7 +85,7 @@ impl ConnectedStream {
         });
 
         match result {
-            Ok((Ok(0), need_poll)) if self.is_receiving_closed.load(Ordering::Relaxed) => {
+            Ok((Ok(0), need_poll)) if self.is_receiving_shut.load(Ordering::Relaxed) => {
                 Ok((0, need_poll))
             }
             Ok((Ok(0), need_poll)) => {
@@ -175,17 +166,9 @@ impl ConnectedStream {
 
     pub(super) fn check_io_events(&self) -> IoEvents {
         self.tcp_conn.raw_with(|socket| {
-            if socket.is_peer_closed() {
+            let is_receiving_closed =
-                // Only the sending side of peer socket is closed
+                self.is_receiving_shut.load(Ordering::Relaxed) || !socket.may_recv_new();
-                self.is_receiving_closed.store(true, Ordering::Relaxed);
+            let is_sending_closed = !socket.may_send();
-            } else if socket.is_closed() {
-                // The sending side of both peer socket and this socket are closed
-                self.is_receiving_closed.store(true, Ordering::Relaxed);
-                self.is_sending_closed.store(true, Ordering::Relaxed);
-            }
-
-            let is_receiving_closed = self.is_receiving_closed.load(Ordering::Relaxed);
-            let is_sending_closed = self.is_sending_closed.load(Ordering::Relaxed);
 
             let mut events = IoEvents::empty();
 
@@ -219,7 +202,7 @@ impl ConnectedStream {
         set_option(&self.tcp_conn)
     }
 
-    pub(super) fn raw_with<R>(&self, f: impl FnOnce(&RawTcpSocket) -> R) -> R {
+    pub(super) fn raw_with<R>(&self, f: impl FnOnce(&RawTcpSocketExt) -> R) -> R {
         self.tcp_conn.raw_with(f)
     }
 }

kernel/src/net/socket/ip/stream/observer.rs

@@ -25,13 +25,18 @@ impl SocketEventObserver for StreamObserver {
             io_events |= IoEvents::OUT;
         }
 
-        if events.contains(SocketEvents::PEER_CLOSED) {
+        if events.contains(SocketEvents::CLOSED_RECV) {
+            // `CLOSED_RECV` definitely causes IN and RDHUP.
             io_events |= IoEvents::IN | IoEvents::RDHUP;
+            // `CLOSED_RECV` may cause HUP/ERR (combined with a previous `CLOSED_SEND`).
+            io_events |= IoEvents::HUP | IoEvents::ERR;
         }
 
-        if events.contains(SocketEvents::CLOSED) {
+        if events.contains(SocketEvents::CLOSED_SEND) {
-            io_events |= IoEvents::IN | IoEvents::OUT;
+            // `CLOSED_SEND` definitely causes OUT.
-            io_events |= IoEvents::RDHUP | IoEvents::HUP | IoEvents::ERR;
+            io_events |= IoEvents::OUT;
+            // `CLOSED_SEND` may cause HUP/ERR (combined with a previous `CLOSED_RECV`).
+            io_events |= IoEvents::HUP | IoEvents::ERR;
         }
 
         self.0.notify(io_events);