Skip to content

Make TCP flag checks more readable #115

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SajjadPourali merged 1 commit into from
Nov 2, 2026
Merged

Make TCP flag checks more readable #115

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 44 additions & 19 deletions src/packet.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,40 @@ impl std::fmt::Display for NetworkTuple {
}
}

pub mod tcp_flags {
pub const CWR: u8 = 0b10000000;
pub const ECE: u8 = 0b01000000;
pub const URG: u8 = 0b00100000;
pub const ACK: u8 = 0b00010000;
pub const PSH: u8 = 0b00001000;
pub const RST: u8 = 0b00000100;
pub const SYN: u8 = 0b00000010;
pub const FIN: u8 = 0b00000001;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct TcpFlags(u8);

impl TcpFlags {
pub(crate) const CWR: Self = Self(0b10000000);
pub(crate) const ECE: Self = Self(0b01000000);
pub(crate) const URG: Self = Self(0b00100000);
pub(crate) const ACK: Self = Self(0b00010000);
pub(crate) const PSH: Self = Self(0b00001000);
pub(crate) const RST: Self = Self(0b00000100);
pub(crate) const SYN: Self = Self(0b00000010);
pub(crate) const FIN: Self = Self(0b00000001);

fn empty() -> Self {
Self(0)
}

pub(crate) fn contains(self, flags: Self) -> bool {
self.0 & flags.0 == flags.0
}
}

impl std::ops::BitOr for TcpFlags {
type Output = Self;

fn bitor(self, rhs: Self) -> Self {
Self(self.0 | rhs.0)
}
}

impl std::ops::BitOrAssign for TcpFlags {
fn bitor_assign(&mut self, rhs: Self) {
self.0 |= rhs.0;
}
}

#[derive(Debug, Clone)]
Expand Down Expand Up @@ -170,31 +195,31 @@ pub fn tcp_header_fmt(header: &TcpHeader) -> String {
)
}

pub fn tcp_header_flags(inner: &TcpHeader) -> u8 {
let mut flags = 0;
pub(crate) fn tcp_header_flags(inner: &TcpHeader) -> TcpFlags {
let mut flags = TcpFlags::empty();
if inner.cwr {
flags |= tcp_flags::CWR;
flags |= TcpFlags::CWR;
}
if inner.ece {
flags |= tcp_flags::ECE;
flags |= TcpFlags::ECE;
}
if inner.urg {
flags |= tcp_flags::URG;
flags |= TcpFlags::URG;
}
if inner.ack {
flags |= tcp_flags::ACK;
flags |= TcpFlags::ACK;
}
if inner.psh {
flags |= tcp_flags::PSH;
flags |= TcpFlags::PSH;
}
if inner.rst {
flags |= tcp_flags::RST;
flags |= TcpFlags::RST;
}
if inner.syn {
flags |= tcp_flags::SYN;
flags |= TcpFlags::SYN;
}
if inner.fin {
flags |= tcp_flags::FIN;
flags |= TcpFlags::FIN;
}

flags
Expand Down
64 changes: 33 additions & 31 deletions src/stream/tcp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,7 @@ use super::seqnum::SeqNum;
use crate::{
PacketReceiver, PacketSender, TTL,
error::IpStackError,
packet::{
IpHeader, NetworkPacket, NetworkTuple, TransportHeader,
tcp_flags::{ACK, FIN, PSH, RST, SYN},
tcp_header_flags, tcp_header_fmt,
},
packet::{IpHeader, NetworkPacket, NetworkTuple, TcpFlags, TransportHeader, tcp_header_flags, tcp_header_fmt},
stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, PacketType, READ_BUFFER_SIZE, READ_CHUNK, RTO, Tcb, TcpState},
};
use atomic_waker::AtomicWaker;
Expand All @@ -33,6 +29,12 @@ const LAST_ACK_MAX_RETRIES: usize = 3;
const LAST_ACK_TIMEOUT: Duration = Duration::from_millis(500);
const TIMEOUT: Duration = Duration::from_secs(60);

const ACK: TcpFlags = TcpFlags::ACK;
const FIN: TcpFlags = TcpFlags::FIN;
const PSH: TcpFlags = TcpFlags::PSH;
const RST: TcpFlags = TcpFlags::RST;
const SYN: TcpFlags = TcpFlags::SYN;

#[non_exhaustive]
#[derive(Debug, Clone)]
/// TCP configuration
Expand Down Expand Up @@ -745,7 +747,7 @@ async fn tcp_main_logic_loop(
break;
}

if flags & RST == RST {
if flags.contains(RST) {
tcb.change_state(TcpState::Closed);
drop(tcb);
write_notify.wake();
Expand Down Expand Up @@ -773,7 +775,7 @@ async fn tcp_main_logic_loop(
}

match state {
TcpState::SynReceived if flags & ACK == ACK => {
TcpState::SynReceived if flags.contains(ACK) => {
if len > 0 {
receive_segment(
&up_packet_sender,
Expand Down Expand Up @@ -888,7 +890,7 @@ async fn tcp_main_logic_loop(
}
}
TcpState::CloseWait => {
if flags & ACK == ACK && tcb.get_inflight_packets_total_len() == 0 {
if flags.contains(ACK) && tcb.get_inflight_packets_total_len() == 0 {
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK | FIN, None, None)?;
tcb.increase_seq();
tcb.change_state(TcpState::LastAck);
Expand All @@ -911,7 +913,7 @@ async fn tcp_main_logic_loop(
write_notify.wake();
}
}
TcpState::LastAck if flags & ACK == ACK => {
TcpState::LastAck if flags.contains(ACK) => {
tcb.change_state(TcpState::Closed);
tokio::spawn(async move {
if let Err(e) = exit_notifier.send(()).await {
Expand All @@ -922,7 +924,7 @@ async fn tcp_main_logic_loop(
log::trace!("{network_tuple} {state:?}: Received final ACK, transitioned to {new_state:?}");
}
TcpState::FinWait1 => {
if flags & (ACK | FIN) == (ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq {
if flags.contains(ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq {
// If the received packet is an ACK with FIN, we need to send an ACK and change state to TimeWait directly, not to FinWait2
tcb.increase_ack();
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
Expand All @@ -931,7 +933,7 @@ async fn tcp_main_logic_loop(
tokio::spawn(task_wait_to_close(tcb_clone.clone(), exit_notifier, network_tuple, config.two_msl));
let new_state = tcb.get_state();
log::trace!("{network_tuple} {state:?}: Final ACK|FIN received too early, transitioned to {new_state:?} directly");
} else if flags & ACK == ACK {
} else if flags.contains(ACK) {
tcb.change_state(TcpState::FinWait2);
if len > 0 {
// if the other side is still sending data, we need to deal with it like PacketStatus::NewPacket
Expand All @@ -954,20 +956,20 @@ async fn tcp_main_logic_loop(
}
}
TcpState::FinWait2 => {
if flags & (ACK | FIN) == (ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq {
if flags.contains(ACK | FIN) && len == 0 && tcb.get_ack() == incoming_seq {
tcb.increase_ack();
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
tcb.change_state(TcpState::TimeWait);
tokio::spawn(task_wait_to_close(tcb_clone.clone(), exit_notifier, network_tuple, config.two_msl));
let new_state = tcb.get_state();
log::trace!("{network_tuple} {state:?}: Received final ACK|FIN, transitioned to {new_state:?}");
} else if flags & ACK == ACK && len == 0 {
} else if flags.contains(ACK) && len == 0 {
// unnormal case, we do nothing here
let l_ack = tcb.get_ack();
if incoming_seq < l_ack {
log::trace!("{network_tuple} {state:?}: Ignoring duplicate ACK, seq {incoming_seq}, expected {l_ack}");
}
} else if flags & ACK == ACK && len > 0 {
} else if flags.contains(ACK) && len > 0 {
if pkt_type == PacketType::KeepAlive {
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
} else {
Expand All @@ -983,7 +985,7 @@ async fn tcp_main_logic_loop(
)?;
write_notify.wake();
}
if flags & FIN == FIN && tcb.get_ack() == incoming_seq + len as u32 {
if flags.contains(FIN) && tcb.get_ack() == incoming_seq + len as u32 {
tcb.increase_ack();
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
tcb.change_state(TcpState::TimeWait);
Expand All @@ -996,7 +998,7 @@ async fn tcp_main_logic_loop(
log::trace!("{network_tuple} {state:?}: Some unnormal case, we do nothing here");
}
}
TcpState::TimeWait if flags & (ACK | FIN) == (ACK | FIN) => {
TcpState::TimeWait if flags.contains(ACK | FIN) => {
write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?;
// wait to timeout, can't call `tcb.change_state(TcpState::Closed);` to change state here
// now we need to wait for the timeout to reach...
Expand All @@ -1008,7 +1010,7 @@ async fn tcp_main_logic_loop(
tcb.update_send_window(tcp_header);
// ACKs with or without data can retire the last outstanding bytes. The
// shutdown waiter may belong to a different task than the writer.
if flags & ACK != 0 && tcb.get_state() == TcpState::Established && tcb.is_inflight_queue_empty() {
if flags.contains(ACK) && tcb.get_state() == TcpState::Established && tcb.is_inflight_queue_empty() {
shutdown.wake();
}
} // end of loop
Expand Down Expand Up @@ -1106,7 +1108,7 @@ pub(crate) fn write_packet_to_device(
tuple: NetworkTuple,
tcb: &Tcb,
options: Option<&Vec<TcpOptions>>,
flags: u8,
flags: TcpFlags,
seq: Option<SeqNum>,
payload: Option<Vec<u8>>,
) -> std::io::Result<usize> {
Expand All @@ -1116,7 +1118,7 @@ pub(crate) fn write_packet_to_device(
// otherwise advertise zero so the peer enters persist mode until the reader frees space.
let recv_window = tcb.get_recv_window();
// RFC 7323 §2.2: the window field of a segment carrying SYN is not scaled.
let advertised = if flags & SYN != 0 {
let advertised = if flags.contains(SYN) {
recv_window
} else {
tcb.get_scaled_recv_window()
Expand All @@ -1133,7 +1135,7 @@ pub(crate) fn write_packet_to_device(
}
// RFC 7323 §2.2: the window scale option travels only on a segment carrying SYN, and only when
// the peer's SYN carried one.
if flags & SYN != 0
if flags.contains(SYN)
&& let Some(shift) = tcb.get_recv_window_shift()
{
tcp_options.push(TcpOptionElement::WindowScale(shift));
Expand Down Expand Up @@ -1161,7 +1163,7 @@ pub(crate) fn create_raw_packet(
src_addr: SocketAddr,
dst_addr: SocketAddr,
calculate_payload_max_len: impl Fn(usize, usize) -> usize,
flags: u8,
flags: TcpFlags,
ttl: u8,
seq: u32,
ack: u32,
Expand All @@ -1171,11 +1173,11 @@ pub(crate) fn create_raw_packet(
) -> std::io::Result<NetworkPacket> {
let mut tcp_header = etherparse::TcpHeader::new(src_addr.port(), dst_addr.port(), seq, win);
tcp_header.acknowledgment_number = ack;
tcp_header.syn = flags & SYN != 0;
tcp_header.ack = flags & ACK != 0;
tcp_header.rst = flags & RST != 0;
tcp_header.fin = flags & FIN != 0;
tcp_header.psh = flags & PSH != 0;
tcp_header.syn = flags.contains(SYN);
tcp_header.ack = flags.contains(ACK);
tcp_header.rst = flags.contains(RST);
tcp_header.fin = flags.contains(FIN);
tcp_header.psh = flags.contains(PSH);

tcp_header.set_options(options).map_err(|e| std::io::Error::new(InvalidInput, e))?;

Expand Down Expand Up @@ -1317,7 +1319,7 @@ mod tests {
.await
}

async fn shutdown_after_ack(same_task: bool, ack_flags: u8) {
async fn shutdown_after_ack(same_task: bool, ack_flags: TcpFlags) {
let (mut stream, mut up_rx, local_seq) = established().await;
let (write_woken, write_waker) = recording_waker();
let mut write_cx = Context::from_waker(&write_waker);
Expand Down Expand Up @@ -1349,7 +1351,7 @@ mod tests {
shutdown_woken.store(false, Ordering::SeqCst);

// The final ACK can also carry peer data; both paths must wake shutdown.
let payload: &[u8] = if ack_flags & PSH != 0 { b"reply" } else { b"" };
let payload: &[u8] = if ack_flags.contains(PSH) { b"reply" } else { b"" };
stream
.stream_sender()
.send(peer_segment(1001, local_seq + 5, ack_flags, u16::MAX, payload))
Expand Down Expand Up @@ -1602,7 +1604,7 @@ mod tests {
}

/// A segment from the peer carrying `payload` at `seq`, with `flags` and a raw `window`.
fn peer_segment(seq: u32, ack: u32, flags: u8, window: u16, payload: &[u8]) -> NetworkPacket {
fn peer_segment(seq: u32, ack: u32, flags: TcpFlags, window: u16, payload: &[u8]) -> NetworkPacket {
let (src, dst) = (PEER.parse().unwrap(), LOCAL.parse().unwrap());
create_raw_packet(src, dst, |_, _| usize::MAX, flags, TTL, seq, ack, window, payload.to_vec(), &[]).unwrap()
}
Expand Down Expand Up @@ -1634,7 +1636,7 @@ mod tests {
.iter()
.map(|p| tcp_header_of(p).acknowledgment_number)
.collect();
assert_eq!(acks, [1001], "flags {flags:#x}: the out-of-order segment got no duplicate ACK");
assert_eq!(acks, [1001], "flags {flags:?}: the out-of-order segment got no duplicate ACK");

// The gap fills: the ACK covers both segments and the reader sees them in order.
stream
Expand All @@ -1646,7 +1648,7 @@ mod tests {
.iter()
.map(|p| tcp_header_of(p).acknowledgment_number)
.collect();
assert_eq!(acks.last(), Some(&1013), "flags {flags:#x}: the buffered segment was not kept");
assert_eq!(acks.last(), Some(&1013), "flags {flags:?}: the buffered segment was not kept");
let mut buf = [0u8; 12];
tokio::io::AsyncReadExt::read_exact(&mut stream, &mut buf).await.unwrap();
assert_eq!(buf[..8], [1; 8]);
Expand Down
Loading