poll_shutdown never sends the FIN when data was still in flight

[Musyg]

Affected: 1.0.1 (crates.io) and main @ 0f95edc - reproduced on both.

Summary

When an application calls shutdown() on an IpStackTcpStream while bytes it
just wrote are still unacknowledged, the FIN is deferred - correctly - but
nothing ever re-polls poll_shutdown once the peer acknowledges them. The FIN
is therefore never sent. The stream stays in Established until the peer gives
up and closes first, or the session times out.

This is the exact shape of tokio::io::copy_bidirectional: when one half
reaches EOF it writes the remaining bytes to the other half and then calls
poll_shutdown on it.

Why it matters

Any protocol that delimits a response by closing the connection hangs. HTTP/1.0
without Content-Length is the common case: the client receives every byte of
the response, acknowledges it, and then blocks in read_to_end until its own
timeout expires - at which point most clients report a timeout and discard what
they already have. The bytes arrived; only the end-of-stream did not.

Measured against a real proxy over a TUN before it was reduced to the
reproduction below: 208 bytes delivered and acknowledged, then 8 seconds of
silence, and the client sent the first FIN.

Sessions also accumulate, since each one stays open until something else
closes it.

Analysis

poll_shutdown (src/stream/tcp.rs:382-411) defers the FIN while the in-flight
queue is non-empty and stores its waker:

Shutdown::None => {
    if is_ready && state == TcpState::Established {
        // ... send_fin_n_change_state_to_fin_wait1(...)
    }
    self.shutdown.lock().unwrap().pending(cx.waker().clone());
    Poll::Pending
}

The acknowledgement that empties that queue is handled here
(src/stream/tcp.rs:741-743 on main):

PacketType::Ack => {
    write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(());
}

Only write_notify is woken. write_notify holds a waker only if poll_write
returned Pending - and in this scenario it did not, because the write
succeeded. shutdown.ready() is called only from the session-teardown paths
(:298, :316, :355, :500). So no waker is woken, poll_shutdown is never
polled again, and the FIN that the very next poll would have sent is never sent.

Reproduction

Self-contained: no TUN device and no privileges. The stack is driven over a
tokio::io::duplex pipe and the program plays the peer by hand.

peer  -> SYN
stack -> SYN|ACK seq=100
peer  -> ACK (established)
   [app] read 3 bytes, answering and closing
peer  -> PSH|ACK "req"
stack -> data seq=101 len=5
peer  -> ACK up to 106 (nothing is in flight any more)

BUG: three seconds and no FIN.

Exits 1 on the bug, 0 when the FIN arrives. Deterministic here: 3 runs out of
3 on 1.0.1, and 3 out of 3 on main @ 0f95edc.

Cargo.toml and src/main.rs

[package]
name = "ipstack-shutdown-repro"
version = "0.1.0"
edition = "2024"

[dependencies]
ipstack = "1.0.1"   # or { path = "../ipstack" } to test main
env_logger = "0.11"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "io-util", "time"] }

//! Minimal reproduction: `poll_shutdown` defers the FIN while data is in
//! flight, and nothing re-polls it once that data is acknowledged.
//!
//! No TUN device and no privileges: the stack is driven over an in-memory
//! duplex pipe, and this program plays the peer by hand.
//!
//! The application side does exactly what `tokio::io::copy_bidirectional` does
//! when the other half reaches EOF: read the request, write the response, then
//! shut the stream down.
//!
//!   peer  -> SYN
//!   stack -> SYN|ACK
//!   peer  -> ACK                      (established)
//!   peer  -> PSH|ACK "req"
//!   stack -> PSH|ACK "hello"          (application wrote, then shutdown())
//!   peer  -> ACK                      (the in-flight queue is now empty)
//!   stack -> FIN|ACK                  <-- expected, never arrives

use std::net::Ipv4Addr;
use std::time::Duration;

use ipstack::{IpStack, IpStackConfig, IpStackStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

const PEER: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 2);
const TARGET: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 3);
const PEER_PORT: u16 = 40000;
const TARGET_PORT: u16 = 80;
const PEER_ISN: u32 = 1000;
const REQUEST: &[u8] = b"req";
const RESPONSE: &[u8] = b"hello";

const FIN: u8 = 0x01;
const SYN: u8 = 0x02;
const PSH: u8 = 0x08;
const ACK: u8 = 0x10;

#[tokio::main]
async fn main() {
    env_logger::init();

    let mut config = IpStackConfig::default();
    config.mtu_unchecked(1500);
    config.packet_information(false);

    let (device, mut wire) = tokio::io::duplex(65536);
    let mut stack = IpStack::new(config, device);

    tokio::spawn(async move {
        while let Ok(stream) = stack.accept().await {
            if let IpStackStream::Tcp(mut tcp) = stream {
                tokio::spawn(async move {
                    let mut buf = [0u8; 64];
                    let n = tcp.read(&mut buf).await.expect("read");
                    println!("   [app] read {n} bytes, answering and closing");
                    tcp.write_all(RESPONSE).await.expect("write");
                    // This is where it hangs: the FIN is deferred because the
                    // bytes above are still in flight, and the ACK that
                    // retires them never wakes this task.
                    tcp.shutdown().await.expect("shutdown");
                    println!("   [app] shutdown() returned");
                });
            }
        }
    });

    let mut buf = Vec::new();

    send(&mut wire, PEER_ISN, 0, SYN, &[]).await;
    println!("peer  -> SYN");

    let synack = expect_packet(&mut wire, &mut buf, Duration::from_secs(2))
        .await
        .expect("no SYN|ACK from the stack");
    assert_eq!(
        synack.flags & (SYN | ACK),
        SYN | ACK,
        "expected SYN|ACK, got flags {:#04x}",
        synack.flags
    );
    println!("stack -> SYN|ACK seq={}", synack.seq);
    let their_isn = synack.seq;

    send(&mut wire, PEER_ISN + 1, their_isn + 1, ACK, &[]).await;
    println!("peer  -> ACK (established)");

    send(&mut wire, PEER_ISN + 1, their_isn + 1, PSH | ACK, REQUEST).await;
    println!("peer  -> PSH|ACK {:?}", String::from_utf8_lossy(REQUEST));

    // The stack may acknowledge the request on its own before the application
    // answers; skip anything that carries no payload.
    let data = loop {
        let p = expect_packet(&mut wire, &mut buf, Duration::from_secs(2))
            .await
            .expect("no response from the stack");
        if !p.payload.is_empty() {
            break p;
        }
    };
    assert_eq!(data.payload, RESPONSE, "expected the response payload");
    println!("stack -> data seq={} len={}", data.seq, data.payload.len());

    let up_to = data.seq + data.payload.len() as u32;
    send(
        &mut wire,
        PEER_ISN + 1 + REQUEST.len() as u32,
        up_to,
        ACK,
        &[],
    )
    .await;
    println!("peer  -> ACK up to {up_to} (nothing is in flight any more)");

    match expect_packet(&mut wire, &mut buf, Duration::from_secs(3)).await {
        Some(p) if p.flags & FIN != 0 => {
            println!("stack -> FIN|ACK: the stream closed as expected");
        }
        Some(p) => {
            println!("BUG: expected FIN, got flags {:#04x}", p.flags);
            std::process::exit(1);
        }
        None => {
            println!();
            println!("BUG: three seconds and no FIN.");
            println!("The application called shutdown() and the peer acknowledged");
            println!("everything, so the in-flight queue is empty; poll_shutdown");
            println!("would send the FIN if it ran once more, but nothing wakes it.");
            std::process::exit(1);
        }
    }
}

/// One packet as it comes out of the stack.
struct Packet {
    flags: u8,
    seq: u32,
    payload: Vec<u8>,
}

/// Reads until one whole IPv4 packet is available, or the budget runs out.
async fn expect_packet(
    wire: &mut tokio::io::DuplexStream,
    buf: &mut Vec<u8>,
    budget: Duration,
) -> Option<Packet> {
    let deadline = tokio::time::Instant::now() + budget;
    loop {
        if let Some(packet) = take_packet(buf) {
            return Some(packet);
        }
        let mut chunk = [0u8; 2048];
        let left = deadline.checked_duration_since(tokio::time::Instant::now())?;
        let n = tokio::time::timeout(left, wire.read(&mut chunk))
            .await
            .ok()?
            .ok()?;
        if n == 0 {
            return None;
        }
        buf.extend_from_slice(&chunk[..n]);
    }
}

/// Takes the first complete IPv4 packet out of the buffer, if there is one.
fn take_packet(buf: &mut Vec<u8>) -> Option<Packet> {
    if buf.len() < 20 {
        return None;
    }
    let total = u16::from_be_bytes([buf[2], buf[3]]) as usize;
    if buf.len() < total {
        return None;
    }
    let packet: Vec<u8> = buf.drain(..total).collect();
    let ihl = ((packet[0] & 0x0f) as usize) * 4;
    let tcp = &packet[ihl..];
    let offset = ((tcp[12] >> 4) as usize) * 4;
    Some(Packet {
        flags: tcp[13],
        seq: u32::from_be_bytes([tcp[4], tcp[5], tcp[6], tcp[7]]),
        payload: tcp[offset..].to_vec(),
    })
}

async fn send(wire: &mut tokio::io::DuplexStream, seq: u32, ack: u32, flags: u8, payload: &[u8]) {
    let packet = build(seq, ack, flags, payload);
    wire.write_all(&packet).await.expect("write to the device");
    // A real TUN hands the stack exactly one packet per read. A duplex pipe is
    // just bytes, so two packets written back to back can be read together and
    // the second one is then dropped. Leave the stack time to take this one.
    tokio::time::sleep(Duration::from_millis(50)).await;
}

/// An IPv4 + TCP packet from the peer to the target, checksums included.
fn build(seq: u32, ack: u32, flags: u8, payload: &[u8]) -> Vec<u8> {
    let tcp_len = 20 + payload.len();

    let mut tcp = Vec::with_capacity(tcp_len);
    tcp.extend_from_slice(&PEER_PORT.to_be_bytes());
    tcp.extend_from_slice(&TARGET_PORT.to_be_bytes());
    tcp.extend_from_slice(&seq.to_be_bytes());
    tcp.extend_from_slice(&ack.to_be_bytes());
    tcp.push(5 << 4);
    tcp.push(flags);
    tcp.extend_from_slice(&65535u16.to_be_bytes());
    tcp.extend_from_slice(&[0, 0]); // checksum, filled in below
    tcp.extend_from_slice(&[0, 0]); // urgent pointer
    tcp.extend_from_slice(payload);

    let mut pseudo = Vec::new();
    pseudo.extend_from_slice(&PEER.octets());
    pseudo.extend_from_slice(&TARGET.octets());
    pseudo.push(0);
    pseudo.push(6);
    pseudo.extend_from_slice(&(tcp_len as u16).to_be_bytes());
    pseudo.extend_from_slice(&tcp);
    let sum = internet_checksum(&pseudo);
    tcp[16..18].copy_from_slice(&sum.to_be_bytes());

    let mut ip = Vec::with_capacity(20 + tcp_len);
    ip.push(0x45);
    ip.push(0);
    ip.extend_from_slice(&((20 + tcp_len) as u16).to_be_bytes());
    ip.extend_from_slice(&[0, 0]); // identification
    ip.extend_from_slice(&[0x40, 0]); // don't fragment
    ip.push(64); // ttl
    ip.push(6); // tcp
    ip.extend_from_slice(&[0, 0]); // checksum, filled in below
    ip.extend_from_slice(&PEER.octets());
    ip.extend_from_slice(&TARGET.octets());
    let sum = internet_checksum(&ip);
    ip[10..12].copy_from_slice(&sum.to_be_bytes());

    ip.extend_from_slice(&tcp);
    ip
}

/// The Internet checksum, RFC 1071.
fn internet_checksum(bytes: &[u8]) -> u16 {
    let mut sum = 0u32;
    let mut i = 0;
    while i + 1 < bytes.len() {
        sum += u32::from(u16::from_be_bytes([bytes[i], bytes[i + 1]]));
        i += 2;
    }
    if i < bytes.len() {
        sum += u32::from(bytes[i]) << 8;
    }
    while sum >> 16 != 0 {
        sum = (sum & 0xffff) + (sum >> 16);
    }
    !(sum as u16)
}

A possible fix

Wake the deferred shutdown when an acknowledgement empties the in-flight queue.
poll_shutdown then runs once more and sends the FIN itself, so the state
machine keeps its single owner of that transition.

@@ impl Shutdown {
         }
         *self = Shutdown::Ready;
     }
+    /// Wake a deferred shutdown WITHOUT completing it, so that `poll_shutdown`
+    /// runs once more and gets its chance to send the FIN.
+    fn wake(&self) {
+        if let Shutdown::Pending(w) = self {
+            w.wake_by_ref();
+        }
+    }
 
     // Just for comparison purpose
     fn fake_clone(&self) -> Shutdown {
@@ impl IpStackTcpStream {
                 read_notify,
                 data_tx,
                 exit_monitor,
+                shutdown.clone(),
             )
             .await;
@@ async fn tcp_main_logic_loop(
     data_tx: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
     mut exit_monitor: tokio::sync::mpsc::Receiver<()>,
+    shutdown: std::sync::Arc<std::sync::Mutex<Shutdown>>,
 ) -> std::io::Result<()> {
@@ async fn tcp_main_logic_loop(
                         PacketType::Ack => {
                             write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(());
+                            // A shutdown deferred by in-flight data is waiting
+                            // for exactly this acknowledgement.
+                            if tcb.get_inflight_packets_total_len() == 0 {
+                                shutdown.lock().unwrap().wake();
+                            }
                         }
                         PacketType::Invalid => {}

With that applied to main, the reproduction prints
stack -> FIN|ACK: the stream closed as expected and exits 0, three runs out
of three. cargo test (10 + 23) and cargo clippy --all-targets stay clean.

I have not opened a PR - happy to, if the approach suits you. There may be a
neater place for the wake, and PacketType::WindowUpdate may deserve the same
treatment; I only verified the path the reproduction exercises.