rusnel 0.9.0

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;
use dashmap::DashMap;
use quinn::{Connection, RecvStream, SendStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::mpsc;
use tracing::{debug, error, info, info_span, warn, Instrument};

use super::remote::{DynamicTarget, HostPort, OpenConn, RemoteRequest};
use super::tcp::{tunnel_tcp_stream, Counters, TunnelHandleOpt};
use super::tunnel::send_open_conn;
use super::udp::{read_datagram, write_datagram};
use anyhow::{anyhow, Result};

/// Max UDP datagram payload SOCKS5 will accept on either direction. IPv4
/// caps a UDP payload at 65 507 bytes; round up so we never need to grow.
const SOCKS_MAX_DATAGRAM: usize = 65_535;

/// Bound on per (src, target) channel capacity for the UDP-over-SOCKS5
/// relay. UDP is unreliable, so on overflow we drop the datagram (with a
/// debug log) instead of blocking the receive loop.
const SOCKS_UDP_CHANNEL_CAPACITY: usize = 4096;

/// How long a per (src, target) UDP-over-SOCKS5 conn may sit idle before
/// we tear down the QUIC stream. Mirrors the static UDP forward timeout so
/// the two paths age out resources in lockstep.
const SOCKS_UDP_IDLE_TIMEOUT: Duration = Duration::from_secs(60);

pub async fn tunnel_socks_client(
    quic_connection: Connection,
    remote: RemoteRequest,
    handle: TunnelHandleOpt,
    tunnel_id: u64,
) -> Result<()> {
    let local_addr = remote.local_socket_addr();
    let listener = TcpListener::bind(local_addr).await?;
    info!(addr = %local_addr, proto = "socks5", "listening");

    let local_counter = AtomicUsize::new(0);

    loop {
        let (mut local_conn, peer) = listener.accept().await?;
        let connection = quic_connection.clone();
        let remote = remote.clone();
        let tunnel_handle = handle.clone();
        let local_id = (local_counter.fetch_add(1, Ordering::Relaxed) + 1) as u64;

        // Fire-and-forget: each accepted SOCKS5 connection runs to completion
        // in its own task so the accept loop never blocks on a slow handshake
        // (regression of #20 §1).
        tokio::spawn(async move {
            let request = match socks_handshake(&mut local_conn).await {
                Ok(r) => r,
                Err(e) => {
                    let span = info_span!("socks5", peer = %peer);
                    let _g = span.enter();
                    warn!(error = %e, "handshake failed");
                    return;
                }
            };

            match request {
                SocksRequest::Connect(target) => {
                    let (send, recv) = match connection.open_bi().await {
                        Ok(stream) => stream,
                        Err(e) => {
                            let span = info_span!("socks5", peer = %peer);
                            let _g = span.enter();
                            error!(error = %e, "failed to open quic stream");
                            return;
                        }
                    };

                    // One admin conn per SOCKS CONNECT, labelled
                    // with the SOCKS client's source plus the target
                    // it asked for so operators can spot which CONNECT
                    // is blocking on what.
                    let conn_guard = tunnel_handle
                        .as_ref()
                        .map(|h| h.open_conn(Some(format!("{peer}=>{target}"))));
                    let conn_id = conn_guard.as_ref().map(|g| g.id()).unwrap_or(local_id);
                    let counters = conn_guard.as_ref().map(|g| g.counters());

                    let span = info_span!(
                        "conn",
                        conn_id,
                        tunnel_id,
                        peer = %peer,
                        target = %target,
                        proto = "socks5/tcp",
                    );
                    async move {
                        info!("conn opened");
                        let started = std::time::Instant::now();
                        let result = start_client_dynamic_tunnel(
                            local_conn,
                            send,
                            recv,
                            tunnel_id,
                            target,
                            counters.clone(),
                        )
                        .await;
                        let dur_ms = started.elapsed().as_millis() as u64;
                        let snap = counters.as_ref().map(|c| c.snapshot());
                        match (&result, snap) {
                            (Ok(()), Some((bytes_in, bytes_out))) => {
                                info!(bytes_in, bytes_out, dur_ms, "conn closed")
                            }
                            (Ok(()), None) => info!(dur_ms, "conn closed"),
                            (Err(e), _) => warn!(dur_ms, error = %e, "conn closed (error)"),
                        }
                        drop(conn_guard);
                    }
                    .instrument(span)
                    .await;
                }
                SocksRequest::UdpAssociate => {
                    let span = info_span!("socks5-udp", peer = %peer);
                    async move {
                        debug!("UDP ASSOCIATE requested");
                        if let Err(e) = handle_socks_udp_associate(
                            connection,
                            local_conn,
                            &remote,
                            tunnel_handle,
                            peer,
                            tunnel_id,
                        )
                        .await
                        {
                            error!(error = %e, "UDP ASSOCIATE failed");
                        }
                    }
                    .instrument(span)
                    .await;
                }
            }
        });
    }
}

async fn start_client_dynamic_tunnel(
    mut socks_conn: TcpStream,
    mut send_channel: SendStream,
    mut recv_channel: RecvStream,
    tunnel_id: u64,
    target: HostPort,
    counters: Counters,
) -> Result<()> {
    send_open_conn(
        &OpenConn {
            tunnel_id,
            dynamic: Some(DynamicTarget::Tcp(target)),
        },
        &mut send_channel,
        &mut recv_channel,
    )
    .await?;

    socks_conn
        .write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
        .await?;

    tunnel_tcp_stream(socks_conn, send_channel, recv_channel, counters).await?;

    Ok(())
}

/// Decoded SOCKS5 client request: either a TCP CONNECT (with the target
/// host:port we'll reach via the tunnel) or a UDP ASSOCIATE (whose target is
/// per-datagram and parsed later from each UDP packet's header).
enum SocksRequest {
    Connect(HostPort),
    UdpAssociate,
}

async fn socks_handshake(conn: &mut TcpStream) -> Result<SocksRequest> {
    let mut buf = [0u8; 256];
    conn.read_exact(&mut buf[..2]).await?;

    if buf[0] != 0x05 {
        return Err(anyhow!("Unsupported SOCKS version: {}", buf[0]));
    }

    let methods_len = buf[1] as usize;
    conn.read_exact(&mut buf[..methods_len]).await?;

    conn.write_all(&[0x05, 0x00]).await?;

    conn.read_exact(&mut buf[..4]).await?;

    let cmd = buf[1];
    let atyp = buf[3];

    let target = read_socks_addr(conn, atyp).await?;

    match cmd {
        0x01 => Ok(SocksRequest::Connect(target)),
        0x03 => Ok(SocksRequest::UdpAssociate),
        _ => {
            // 0x07 = command not supported
            conn.write_all(&[0x05, 0x07]).await?;
            Err(anyhow!("Unsupported SOCKS command: {}", cmd))
        }
    }
}

/// Read the address portion (`ATYP DST.ADDR DST.PORT`) of a SOCKS5 request
/// or UDP datagram. `atyp` is the already-consumed address-type byte.
async fn read_socks_addr(conn: &mut TcpStream, atyp: u8) -> Result<HostPort> {
    match atyp {
        0x01 => {
            let mut addr = [0u8; 4];
            conn.read_exact(&mut addr).await?;
            let mut port = [0u8; 2];
            conn.read_exact(&mut port).await?;
            Ok(HostPort::new(
                Ipv4Addr::from(addr).to_string(),
                u16::from_be_bytes(port),
            ))
        }
        0x03 => {
            let mut len = [0u8; 1];
            conn.read_exact(&mut len).await?;
            let mut domain = vec![0u8; len[0] as usize];
            conn.read_exact(&mut domain).await?;
            let mut port = [0u8; 2];
            conn.read_exact(&mut port).await?;
            Ok(HostPort::new(
                String::from_utf8_lossy(&domain).into_owned(),
                u16::from_be_bytes(port),
            ))
        }
        0x04 => {
            let mut addr = [0u8; 16];
            conn.read_exact(&mut addr).await?;
            let mut port = [0u8; 2];
            conn.read_exact(&mut port).await?;
            Ok(HostPort::new(
                Ipv6Addr::from(addr).to_string(),
                u16::from_be_bytes(port),
            ))
        }
        other => {
            conn.write_all(&[0x05, 0x08]).await?;
            Err(anyhow!("Unsupported address type: {}", other))
        }
    }
}

// ---------------------------------------------------------------------------
// SOCKS5 UDP ASSOCIATE
// ---------------------------------------------------------------------------
//
// RFC 1928 §6 describes UDP ASSOCIATE: the SOCKS server binds a UDP socket
// and replies with that socket's address. The SOCKS client then sends UDP
// datagrams to that address with a per-datagram header identifying the
// ultimate target. Each datagram looks like:
//
//   +----+------+------+----------+----------+----------+
//   |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
//   +----+------+------+----------+----------+----------+
//   | 2  |  1   |  1   | Variable |    2     | Variable |
//   +----+------+------+----------+----------+----------+
//
// The TCP control connection must stay open: when it closes, the UDP
// association is torn down. We support FRAG=0 only (no reassembly).

async fn handle_socks_udp_associate(
    quic_connection: Connection,
    mut tcp_conn: TcpStream,
    original_remote: &RemoteRequest,
    handle: TunnelHandleOpt,
    socks_peer: SocketAddr,
    tunnel_id: u64,
) -> Result<()> {
    // Bind a UDP socket on the same local IP we listen on. Port 0 lets the
    // OS pick — we report the chosen port back to the SOCKS client in the
    // reply. Loopback/127.0.0.1 is the conventional default and matches
    // what Chisel and most other SOCKS5 servers do.
    let listen_ip = original_remote.local_socket_addr().ip();
    let bind_addr = SocketAddr::new(listen_ip, 0);
    let udp_socket = Arc::new(UdpSocket::bind(bind_addr).await?);
    let bound = udp_socket.local_addr()?;
    debug!(addr = %bound, "UDP ASSOCIATE relay bound");

    write_socks_udp_associate_reply(&mut tcp_conn, bound).await?;

    // Spawn the relay; abort it when the TCP control connection closes.
    let relay_handle = tokio::spawn({
        let connection = quic_connection.clone();
        let socket = udp_socket.clone();
        let handle = handle.clone();
        async move {
            if let Err(e) =
                run_socks_udp_relay(connection, tunnel_id, socket, handle, socks_peer).await
            {
                debug!(error = %e, "UDP relay ended");
            }
        }
    });

    // RFC 1928: the TCP control connection signals the lifetime of the
    // association. It carries no application data — any reads here are
    // either keep-alive bytes or EOF. We drain and wait for the close.
    let mut buf = [0u8; 64];
    while let Ok(n) = tcp_conn.read(&mut buf).await {
        if n == 0 {
            break;
        }
    }
    relay_handle.abort();
    debug!("UDP ASSOCIATE control closed");
    Ok(())
}

/// Write the SOCKS5 reply for UDP ASSOCIATE: VER=5 REP=00 RSV=00 ATYP
/// BND.ADDR BND.PORT — where BND.ADDR/PORT is our UDP relay socket.
async fn write_socks_udp_associate_reply(conn: &mut TcpStream, addr: SocketAddr) -> Result<()> {
    let mut reply = vec![0x05, 0x00, 0x00];
    match addr.ip() {
        IpAddr::V4(v4) => {
            reply.push(0x01);
            reply.extend_from_slice(&v4.octets());
        }
        IpAddr::V6(v6) => {
            reply.push(0x04);
            reply.extend_from_slice(&v6.octets());
        }
    }
    reply.extend_from_slice(&addr.port().to_be_bytes());
    conn.write_all(&reply).await?;
    Ok(())
}

/// Per-(source, target) UDP-over-SOCKS5 relay loop. Reads datagrams from the
/// SOCKS-bound UDP socket, parses the SOCKS5 UDP header to discover the
/// ultimate target, and ferries the inner payload through a per-target QUIC
/// stream. Replies arriving on each QUIC stream are wrapped back into SOCKS5
/// UDP framing and sent to the original SOCKS UDP source.
async fn run_socks_udp_relay(
    quic_connection: Connection,
    tunnel_id: u64,
    udp_socket: Arc<UdpSocket>,
    handle: TunnelHandleOpt,
    socks_peer: SocketAddr,
) -> Result<()> {
    // Key by (source, target): this preserves source isolation (so two
    // different SOCKS clients sharing the same association never see each
    // other's traffic) and target isolation (so replies on each per-target
    // QUIC stream can be wrapped with the correct DST.ADDR/PORT header).
    let conns: Arc<DashMap<(SocketAddr, HostPort), mpsc::Sender<Bytes>>> = Arc::new(DashMap::new());

    let mut buf = vec![0u8; SOCKS_MAX_DATAGRAM];
    loop {
        let (n, src) = udp_socket.recv_from(&mut buf).await?;
        let (target, payload) = match parse_socks_udp_header(&buf[..n]) {
            Ok(v) => v,
            Err(e) => {
                debug!(peer = %src, error = %e, "invalid UDP datagram");
                continue;
            }
        };

        let key = (src, target.clone());
        let mut existing = conns.get(&key).map(|e| e.value().clone());
        if let Some(tx) = &existing {
            if tx.is_closed() {
                conns.remove(&key);
                existing = None;
            }
        }
        let tx = match existing {
            Some(tx) => tx,
            None => {
                let (tx, rx) = mpsc::channel(SOCKS_UDP_CHANNEL_CAPACITY);
                conns.insert(key.clone(), tx.clone());
                spawn_socks_udp_conn(
                    quic_connection.clone(),
                    tunnel_id,
                    udp_socket.clone(),
                    src,
                    target.clone(),
                    rx,
                    conns.clone(),
                    handle.clone(),
                    socks_peer,
                );
                tx
            }
        };

        // UDP is unreliable: drop on backpressure or after the conn
        // terminated rather than blocking the receive loop.
        if let Err(e) = tx.try_send(Bytes::copy_from_slice(payload)) {
            debug!(peer = %src, target = %target, error = %e, "dropping udp datagram");
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn spawn_socks_udp_conn(
    quic_connection: Connection,
    tunnel_id: u64,
    udp_socket: Arc<UdpSocket>,
    source: SocketAddr,
    target: HostPort,
    rx: mpsc::Receiver<Bytes>,
    conns: Arc<DashMap<(SocketAddr, HostPort), mpsc::Sender<Bytes>>>,
    handle: TunnelHandleOpt,
    socks_peer: SocketAddr,
) {
    let key = (source, target.clone());
    tokio::spawn(async move {
        // One admin conn per (SOCKS-source, UDP-target) pair so the
        // operator can spot a single greedy target over a shared
        // SOCKS UDP association.
        let conn_guard = handle
            .as_ref()
            .map(|h| h.open_conn(Some(format!("{socks_peer}=>{target}"))));
        let conn_id = conn_guard.as_ref().map(|g| g.id()).unwrap_or(0);
        let counters = conn_guard.as_ref().map(|g| g.counters());
        let span = info_span!(
            "conn",
            conn_id,
            tunnel_id,
            peer = %socks_peer,
            target = %target,
            proto = "socks5/udp",
        );
        async move {
            info!("conn opened");
            let started = std::time::Instant::now();
            let result = run_socks_udp_conn(
                quic_connection,
                tunnel_id,
                udp_socket,
                source,
                target,
                rx,
                counters.clone(),
            )
            .await;
            let dur_ms = started.elapsed().as_millis() as u64;
            let snap = counters.as_ref().map(|c| c.snapshot());
            match (&result, snap) {
                (Ok(()), Some((bytes_in, bytes_out))) => {
                    info!(bytes_in, bytes_out, dur_ms, "conn closed")
                }
                (Ok(()), None) => info!(dur_ms, "conn closed"),
                (Err(e), _) => warn!(dur_ms, error = %e, "conn closed (error)"),
            }
        }
        .instrument(span)
        .await;
        drop(conn_guard);
        conns.remove(&key);
    });
}

#[allow(clippy::too_many_arguments)]
async fn run_socks_udp_conn(
    quic_connection: Connection,
    tunnel_id: u64,
    udp_socket: Arc<UdpSocket>,
    source: SocketAddr,
    target: HostPort,
    mut rx: mpsc::Receiver<Bytes>,
    counters: Counters,
) -> Result<()> {
    // Open a fresh QUIC bi-stream and announce it as a SOCKS5 dynamic
    // UDP conn under the parent SOCKS5 tunnel — the server uses the
    // attached `DynamicTarget::Udp(target)` to dial the right
    // destination. Same length-prefixed datagram framing as the
    // static UDP forward path beyond this point.
    let (mut send_channel, mut recv_channel) = quic_connection.open_bi().await?;
    send_open_conn(
        &OpenConn {
            tunnel_id,
            dynamic: Some(DynamicTarget::Udp(target.clone())),
        },
        &mut send_channel,
        &mut recv_channel,
    )
    .await?;

    // Forward SOCKS-side datagrams onto the QUIC stream until the rx side
    // closes (relay loop dropped the conn) or the conn sits idle long
    // enough to age out.
    let local_to_quic = async {
        loop {
            match tokio::time::timeout(SOCKS_UDP_IDLE_TIMEOUT, rx.recv()).await {
                Ok(Some(payload)) => {
                    write_datagram(&mut send_channel, &payload).await?;
                    if let Some(c) = counters.as_ref() {
                        c.add_out(payload.len() as u64);
                    }
                }
                Ok(None) => return Ok::<(), anyhow::Error>(()),
                Err(_) => {
                    debug!("idle timeout");
                    return Ok(());
                }
            }
        }
    };

    // Wrap each QUIC-arriving reply in a SOCKS5 UDP header naming this
    // conn's target, then send it to the original SOCKS client UDP
    // source.
    let quic_to_local = async {
        let mut buf = vec![0u8; SOCKS_MAX_DATAGRAM];
        let mut wrap = Vec::with_capacity(SOCKS_MAX_DATAGRAM + 32);
        loop {
            let payload = read_datagram(&mut recv_channel, &mut buf).await?;
            wrap_socks_udp_reply(&target, payload, &mut wrap);
            udp_socket.send_to(&wrap, source).await?;
            if let Some(c) = counters.as_ref() {
                c.add_in(payload.len() as u64);
            }
        }
    };

    tokio::select! {
        r = local_to_quic => r,
        r = quic_to_local => r,
    }
}

/// Parse a SOCKS5 UDP datagram header, returning the target the datagram is
/// destined for and the inner payload slice. Rejects fragmented datagrams
/// (FRAG != 0) — the spec allows but does not require their support.
fn parse_socks_udp_header(buf: &[u8]) -> Result<(HostPort, &[u8])> {
    if buf.len() < 4 {
        return Err(anyhow!("UDP header too short"));
    }
    if buf[2] != 0 {
        return Err(anyhow!("UDP fragmentation not supported"));
    }
    let atyp = buf[3];
    let (target, hdr_len) = match atyp {
        0x01 => {
            if buf.len() < 4 + 4 + 2 {
                return Err(anyhow!("truncated IPv4 header"));
            }
            let ip = Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
            let port = u16::from_be_bytes([buf[8], buf[9]]);
            (HostPort::new(ip.to_string(), port), 10)
        }
        0x04 => {
            if buf.len() < 4 + 16 + 2 {
                return Err(anyhow!("truncated IPv6 header"));
            }
            let mut octets = [0u8; 16];
            octets.copy_from_slice(&buf[4..20]);
            let ip = Ipv6Addr::from(octets);
            let port = u16::from_be_bytes([buf[20], buf[21]]);
            (HostPort::new(ip.to_string(), port), 22)
        }
        0x03 => {
            if buf.len() < 5 {
                return Err(anyhow!("truncated domain header"));
            }
            let len = buf[4] as usize;
            if buf.len() < 5 + len + 2 {
                return Err(anyhow!("truncated domain header"));
            }
            let domain = String::from_utf8_lossy(&buf[5..5 + len]).into_owned();
            let port = u16::from_be_bytes([buf[5 + len], buf[5 + len + 1]]);
            (HostPort::new(domain, port), 5 + len + 2)
        }
        other => return Err(anyhow!("unknown ATYP: {}", other)),
    };
    Ok((target, &buf[hdr_len..]))
}

/// Wrap a reply UDP payload with the SOCKS5 UDP datagram header, addressed
/// from the target this conn is talking to. Re-uses the supplied scratch
/// buffer so we don't allocate per reply.
fn wrap_socks_udp_reply(target: &HostPort, payload: &[u8], buf: &mut Vec<u8>) {
    buf.clear();
    buf.extend_from_slice(&[0, 0, 0]);
    if let Ok(ip) = target.host.parse::<Ipv4Addr>() {
        buf.push(0x01);
        buf.extend_from_slice(&ip.octets());
    } else if let Ok(ip) = target.host.parse::<Ipv6Addr>() {
        buf.push(0x04);
        buf.extend_from_slice(&ip.octets());
    } else {
        buf.push(0x03);
        let bytes = target.host.as_bytes();
        // Domain length is bounded by `u8::MAX`; longer names can't be
        // represented in the SOCKS5 wire format, so truncate defensively.
        let len = bytes.len().min(u8::MAX as usize);
        buf.push(len as u8);
        buf.extend_from_slice(&bytes[..len]);
    }
    buf.extend_from_slice(&target.port.to_be_bytes());
    buf.extend_from_slice(payload);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_udp_header_ipv4() {
        // RSV RSV FRAG ATYP=1 IP=1.2.3.4 PORT=53 DATA="abc"
        let buf = [0, 0, 0, 1, 1, 2, 3, 4, 0, 53, b'a', b'b', b'c'];
        let (target, payload) = parse_socks_udp_header(&buf).unwrap();
        assert_eq!(target.host, "1.2.3.4");
        assert_eq!(target.port, 53);
        assert_eq!(payload, b"abc");
    }

    #[test]
    fn parse_udp_header_domain() {
        let mut buf = vec![0, 0, 0, 3, 11];
        buf.extend_from_slice(b"example.com");
        buf.extend_from_slice(&80u16.to_be_bytes());
        buf.extend_from_slice(b"hi");
        let (target, payload) = parse_socks_udp_header(&buf).unwrap();
        assert_eq!(target.host, "example.com");
        assert_eq!(target.port, 80);
        assert_eq!(payload, b"hi");
    }

    #[test]
    fn parse_udp_header_ipv6() {
        let mut buf = vec![0, 0, 0, 4];
        buf.extend_from_slice(&Ipv6Addr::LOCALHOST.octets());
        buf.extend_from_slice(&443u16.to_be_bytes());
        buf.extend_from_slice(b"x");
        let (target, payload) = parse_socks_udp_header(&buf).unwrap();
        assert_eq!(target.host, "::1");
        assert_eq!(target.port, 443);
        assert_eq!(payload, b"x");
    }

    #[test]
    fn parse_udp_header_rejects_fragmentation() {
        let buf = [0, 0, 1, 1, 1, 2, 3, 4, 0, 53];
        assert!(parse_socks_udp_header(&buf).is_err());
    }

    #[test]
    fn wrap_reply_ipv4_roundtrip() {
        let target = HostPort::new("8.8.8.8", 53);
        let mut out = Vec::new();
        wrap_socks_udp_reply(&target, b"hello", &mut out);
        let (parsed, payload) = parse_socks_udp_header(&out).unwrap();
        assert_eq!(parsed, target);
        assert_eq!(payload, b"hello");
    }

    #[test]
    fn wrap_reply_domain_roundtrip() {
        let target = HostPort::new("example.com", 80);
        let mut out = Vec::new();
        wrap_socks_udp_reply(&target, b"data", &mut out);
        let (parsed, payload) = parse_socks_udp_header(&out).unwrap();
        assert_eq!(parsed, target);
        assert_eq!(payload, b"data");
    }
}