Existing connections keep sending GSO batches after noq-udp disables GSO #4555

//! Repro: iroh keeps sending GSO batches on a connection after noq-udp has disabled GSO.
//!
//! On the affected machine (Linux before 6.11, NIC without TX checksum offload):
//!     cargo run --release -- server 4433
//! On any other machine:
//!     cargo run --release -- client <server-ip>:4433
//!
//! Passing `--no-gso` to the server turns segmentation offload off, which avoids the bug.

use std::{error::Error, net::ToSocketAddrs, time::Instant};

use iroh::{
    Endpoint, EndpointAddr, RelayMode, SecretKey,
    endpoint::{Incoming, QuicTransportConfig, presets},
};

const ALPN: &[u8] = b"gso-repro";
const REPLY_SIZE: usize = 256 * 1024;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    tracing_subscriber::fmt()
        .with_env_filter("noq_udp=info")
        .init();

    let args: Vec<String> = std::env::args().collect();
    match args.get(1).map(String::as_str) {
        Some("server") => {
            let port: u16 = args[2].parse()?;
            let use_gso = !args.iter().any(|arg| arg == "--no-gso");
            run_server(port, use_gso).await
        }
        Some("client") => run_client(&args[2]).await,
        _ => Err("usage: server <port> [--no-gso] | client <server-ip>:<port>".into()),
    }
}

/// Both sides use this key, so the client knows the server's id without copying it over.
fn server_secret_key() -> SecretKey {
    SecretKey::from_bytes(&[7; 32])
}

async fn run_server(port: u16, use_gso: bool) -> Result<(), Box<dyn Error>> {
    let transport = QuicTransportConfig::builder()
        .enable_segmentation_offload(use_gso)
        .build();
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(server_secret_key())
        .relay_mode(RelayMode::Disabled)
        .bind_addr(format!("0.0.0.0:{port}"))?
        .alpns(vec![ALPN.to_vec()])
        .transport_config(transport)
        .bind()
        .await?;
    println!("listening on udp port {port}, gso {use_gso}");

    while let Some(incoming) = endpoint.accept().await {
        if let Err(err) = serve_one_request(incoming).await {
            println!("connection failed: {err}");
        }
    }
    Ok(())
}

/// Answers one request with `REPLY_SIZE` bytes, then waits for the client to close.
async fn serve_one_request(incoming: Incoming) -> Result<(), Box<dyn Error>> {
    let connection = incoming.await?;
    let (mut send, mut recv) = connection.accept_bi().await?;

    let mut request = [0u8; 1];
    recv.read_exact(&mut request).await?;
    send.write_all(&vec![0u8; REPLY_SIZE]).await?;

    let reason = connection.closed().await;
    println!("connection closed: {reason}");
    Ok(())
}

async fn run_client(server: &str) -> Result<(), Box<dyn Error>> {
    let server_ip = server
        .to_socket_addrs()?
        .next()
        .ok_or("cannot resolve the server")?;
    let server_addr = EndpointAddr::new(server_secret_key().public()).with_ip_addr(server_ip);

    let endpoint = Endpoint::builder(presets::Minimal)
        .relay_mode(RelayMode::Disabled)
        .bind()
        .await?;
    let connection = endpoint.connect(server_addr, ALPN).await?;
    let (mut send, mut recv) = connection.open_bi().await?;

    let started = Instant::now();
    send.write_all(&[1]).await?;
    let mut reply = vec![0u8; REPLY_SIZE];
    let result = recv.read_exact(&mut reply).await;
    println!(
        "{} KiB reply after {:.1?}: {result:?}",
        REPLY_SIZE / 1024,
        started.elapsed()
    );
    result?;

    connection.close(0u32.into(), b"done");
    endpoint.close().await;
    Ok(())
}