From eef56c770b60da39183cebb913f471ef7da5f214 Mon Sep 17 00:00:00 2001 From: Ruihan Li Date: Wed, 19 Feb 2025 22:17:33 +0800 Subject: [PATCH] Simplify the TCP state check --- .../libs/aster-bigtcp/src/socket/bound/mod.rs | 2 +- .../aster-bigtcp/src/socket/bound/tcp_conn.rs | 46 +++++++++++++------ .../src/socket/bound/tcp_listen.rs | 3 +- kernel/libs/aster-bigtcp/src/socket/event.rs | 6 ++- kernel/libs/aster-bigtcp/src/socket/mod.rs | 13 +++--- kernel/libs/aster-bigtcp/src/socket/option.rs | 2 +- kernel/libs/aster-bigtcp/src/socket/state.rs | 31 ------------- .../libs/aster-bigtcp/src/socket/unbound.rs | 3 +- kernel/src/net/iface/mod.rs | 2 + kernel/src/net/socket/ip/stream/connected.rs | 39 +++++----------- kernel/src/net/socket/ip/stream/observer.rs | 13 ++++-- 11 files changed, 70 insertions(+), 90 deletions(-) delete mode 100644 kernel/libs/aster-bigtcp/src/socket/state.rs diff --git a/kernel/libs/aster-bigtcp/src/socket/bound/mod.rs b/kernel/libs/aster-bigtcp/src/socket/bound/mod.rs index ccf49306f..9045085c0 100644 --- a/kernel/libs/aster-bigtcp/src/socket/bound/mod.rs +++ b/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; diff --git a/kernel/libs/aster-bigtcp/src/socket/bound/tcp_conn.rs b/kernel/libs/aster-bigtcp/src/socket/bound/tcp_conn.rs index 8987b4d88..c18f190a9 100644 --- a/kernel/libs/aster-bigtcp/src/socket/bound/tcp_conn.rs +++ b/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, - RawTcpSocket, TcpStateCheck, + unbound::{new_tcp_socket, RawTcpSocket}, }, socket_table::ConnectionKey, }; @@ -37,7 +36,7 @@ pub struct TcpConnectionInner { connection_key: ConnectionKey, } -pub(super) struct RawTcpSocketExt { +pub struct RawTcpSocketExt { socket: Box, pub(super) listener: Option>>, has_connected: bool, @@ -57,6 +56,25 @@ impl DerefMut for RawTcpSocketExt { } } +impl RawTcpSocketExt { + /// 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 RawTcpSocketExt { &mut self, this: &Arc>, ) -> (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 RawTcpSocketExt { let became_dead = self.check_dead(this); - let events = if self.is_peer_closed() { - SocketEvents::PEER_CLOSED - } else if self.is_closed() { - SocketEvents::CLOSED - } else { - SocketEvents::empty() - }; + let mut events = SocketEvents::empty(); + if !self.may_recv_new() { + events |= SocketEvents::CLOSED_RECV; + } + if !may_send { + events |= SocketEvents::CLOSED_SEND; + } (events, became_dead) } @@ -317,7 +337,7 @@ impl TcpConnection { socket.listener = None; - if socket.is_closed() { + if socket.state() == State::Closed { return false; } @@ -333,7 +353,7 @@ impl TcpConnection { // polling time. pub fn raw_with(&self, f: F) -> R where - F: FnOnce(&RawTcpSocket) -> R, + F: FnOnce(&RawTcpSocketExt) -> R, { let socket = self.0.inner.lock(); f(&socket) diff --git a/kernel/libs/aster-bigtcp/src/socket/bound/tcp_listen.rs b/kernel/libs/aster-bigtcp/src/socket/bound/tcp_listen.rs index bbb72bd2e..195052ee9 100644 --- a/kernel/libs/aster-bigtcp/src/socket/bound/tcp_listen.rs +++ b/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, - RawTcpSocket, + unbound::{new_tcp_socket, RawTcpSocket}, }, socket_table::{ConnectionKey, ListenerKey}, }; diff --git a/kernel/libs/aster-bigtcp/src/socket/event.rs b/kernel/libs/aster-bigtcp/src/socket/event.rs index 0aba28c26..c51514c29 100644 --- a/kernel/libs/aster-bigtcp/src/socket/event.rs +++ b/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; - const CLOSED = 8; + /// Receiving new data isn't possible anymore. + const CLOSED_RECV = 4; + /// Sending data isn't possible anymore. + const CLOSED_SEND = 8; } } diff --git a/kernel/libs/aster-bigtcp/src/socket/mod.rs b/kernel/libs/aster-bigtcp/src/socket/mod.rs index 136a635a9..78ff3c925 100644 --- a/kernel/libs/aster-bigtcp/src/socket/mod.rs +++ b/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::{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>; +pub use unbound::{ + RawUdpSocket, TCP_RECV_BUF_LEN, TCP_SEND_BUF_LEN, UDP_RECV_PAYLOAD_LEN, UDP_SEND_PAYLOAD_LEN, +}; diff --git a/kernel/libs/aster-bigtcp/src/socket/option.rs b/kernel/libs/aster-bigtcp/src/socket/option.rs index e597a5ba8..b965c7aa5 100644 --- a/kernel/libs/aster-bigtcp/src/socket/option.rs +++ b/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 { diff --git a/kernel/libs/aster-bigtcp/src/socket/state.rs b/kernel/libs/aster-bigtcp/src/socket/state.rs deleted file mode 100644 index f8bc4148e..000000000 --- a/kernel/libs/aster-bigtcp/src/socket/state.rs +++ /dev/null @@ -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 - } -} diff --git a/kernel/libs/aster-bigtcp/src/socket/unbound.rs b/kernel/libs/aster-bigtcp/src/socket/unbound.rs index 724298c80..730dcf2b4 100644 --- a/kernel/libs/aster-bigtcp/src/socket/unbound.rs +++ b/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 { let raw_tcp_socket = { diff --git a/kernel/src/net/iface/mod.rs b/kernel/src/net/iface/mod.rs index 6284b015f..504738aea 100644 --- a/kernel/src/net/iface/mod.rs +++ b/kernel/src/net/iface/mod.rs @@ -11,6 +11,8 @@ pub use poll::lazy_init; pub type Iface = dyn aster_bigtcp::iface::Iface; pub type BoundPort = aster_bigtcp::iface::BoundPort; +pub type RawTcpSocketExt = aster_bigtcp::socket::RawTcpSocketExt; + pub type TcpConnection = aster_bigtcp::socket::TcpConnection; pub type TcpListener = aster_bigtcp::socket::TcpListener; pub type UdpSocket = aster_bigtcp::socket::UdpSocket; diff --git a/kernel/src/net/socket/ip/stream/connected.rs b/kernel/src/net/socket/ip/stream/connected.rs index fa98c7d31..9cf9feb4d 100644 --- a/kernel/src/net/socket/ip/stream/connected.rs +++ b/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. - /// - /// 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, + /// Indicates if the receiving side of this socket is shut down by the user. + is_receiving_shut: AtomicBool, } impl ConnectedStream { @@ -55,8 +48,7 @@ impl ConnectedStream { tcp_conn, remote_endpoint, is_new_connection, - is_receiving_closed: AtomicBool::new(false), - is_sending_closed: AtomicBool::new(false), + is_receiving_shut: 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() { - // Only the sending side of peer socket is closed - self.is_receiving_closed.store(true, Ordering::Relaxed); - } 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 is_receiving_closed = + self.is_receiving_shut.load(Ordering::Relaxed) || !socket.may_recv_new(); + let is_sending_closed = !socket.may_send(); let mut events = IoEvents::empty(); @@ -219,7 +202,7 @@ impl ConnectedStream { set_option(&self.tcp_conn) } - pub(super) fn raw_with(&self, f: impl FnOnce(&RawTcpSocket) -> R) -> R { + pub(super) fn raw_with(&self, f: impl FnOnce(&RawTcpSocketExt) -> R) -> R { self.tcp_conn.raw_with(f) } } diff --git a/kernel/src/net/socket/ip/stream/observer.rs b/kernel/src/net/socket/ip/stream/observer.rs index cb27d9180..214828cc2 100644 --- a/kernel/src/net/socket/ip/stream/observer.rs +++ b/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) { - io_events |= IoEvents::IN | IoEvents::OUT; - io_events |= IoEvents::RDHUP | IoEvents::HUP | IoEvents::ERR; + if events.contains(SocketEvents::CLOSED_SEND) { + // `CLOSED_SEND` definitely causes OUT. + 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);