Simplify the TCP state check

This commit is contained in:
Ruihan Li
2025-02-19 22:17:33 +08:00
committed by Tate, Hongliang Tian
parent 9edee83ef1
commit eef56c770b
11 changed files with 70 additions and 90 deletions

View File

@ -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;

View File

@ -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<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() {
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<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)

View File

@ -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},
};

View File

@ -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;
}
}

View File

@ -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,
};

View File

@ -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 {

View File

@ -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
}
}

View File

@ -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 = {

View File

@ -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>;

View File

@ -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<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)
}
}

View File

@ -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);