SoulStream/crates/soulseek-client/src/conn.rs

//! Connection plumbing: server socket IO, peer listener, peer message
//! loops, indirect (PierceFirewall) connections, and transfer-connection
//! establishment in all four direct/indirect × inbound/outbound flavors.

use bytes::Bytes;
use std::net::SocketAddrV4;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, oneshot};

use soulseek_proto::{peer, server, MAX_MESSAGE_SIZE};

use crate::{ClientError, Event, Inner, PeerHandle, Result, TransferAck};

pub(crate) const PEER_CONNECT_TIMEOUT: Duration = Duration::from_secs(8);
pub(crate) const INDIRECT_TIMEOUT: Duration = Duration::from_secs(15);
/// Idle cutoff for peer message connections.
const PEER_READ_TIMEOUT: Duration = Duration::from_secs(60);
const SERVER_PING_INTERVAL: Duration = Duration::from_secs(300);

/// Read one length-prefixed frame; returns the frame body (code + payload).
pub(crate) async fn read_frame<R: AsyncReadExt + Unpin>(
    stream: &mut R,
) -> std::io::Result<Option<Bytes>> {
    let mut len_buf = [0u8; 4];
    match stream.read_exact(&mut len_buf).await {
        Ok(_) => {}
        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(e) => return Err(e),
    }
    let len = u32::from_le_bytes(len_buf) as usize;
    if len == 0 || len > MAX_MESSAGE_SIZE {
        return Err(std::io::Error::other(format!("bad frame length {len}")));
    }
    let mut body = vec![0u8; len];
    stream.read_exact(&mut body).await?;
    Ok(Some(Bytes::from(body)))
}

pub(crate) async fn bind_listener(preferred_port: u16) -> std::io::Result<TcpListener> {
    // Walk forward a few ports if the preferred one is taken.
    for offset in 0..10 {
        match TcpListener::bind(("0.0.0.0", preferred_port + offset)).await {
            Ok(l) => return Ok(l),
            Err(_) if offset < 9 => continue,
            Err(e) => return Err(e),
        }
    }
    unreachable!()
}

// ---------------------------------------------------------------- server IO

pub(crate) fn spawn_server_io(
    stream: TcpStream,
    mut outbound: mpsc::UnboundedReceiver<Vec<u8>>,
    mut shutdown: mpsc::UnboundedReceiver<()>,
    inner: Arc<Inner>,
) {
    tokio::spawn(async move {
        let (mut rd, mut wr) = stream.into_split();

        let writer = tokio::spawn(async move {
            while let Some(frame) = outbound.recv().await {
                if wr.write_all(&frame).await.is_err() {
                    break;
                }
            }
        });

        let ping_inner = inner.clone();
        let ping = tokio::spawn(async move {
            loop {
                tokio::time::sleep(SERVER_PING_INTERVAL).await;
                if ping_inner.send_server(server::ping()).is_err() {
                    break;
                }
            }
        });

        let reason = loop {
            let frame = tokio::select! {
                f = read_frame(&mut rd) => f,
                _ = shutdown.recv() => break "shutdown requested".to_string(),
            };
            match frame {
                Ok(Some(body)) => handle_server_frame(&inner, body),
                Ok(None) => break "server closed the connection".to_string(),
                Err(e) => break format!("server read error: {e}"),
            }
        };

        writer.abort();
        ping.abort();
        let _ = inner.event_tx.send(Event::Disconnected { reason });
    });
}

fn handle_server_frame(inner: &Arc<Inner>, mut body: Bytes) {
    use bytes::Buf;
    if body.len() < 4 {
        return;
    }
    let code = body.get_u32_le();

    let event = match server::decode(code, body) {
        Ok(ev) => ev,
        Err(e) => {
            log::warn!("failed to decode server message {code}: {e}");
            return;
        }
    };

    match event {
        server::ServerEvent::LoginOk {
            greeting, own_ip, ..
        } => {
            if let Some(tx) = inner.pending_login.lock().unwrap().take() {
                let _ = tx.send(Ok(greeting.clone()));
            }
            let _ = inner.event_tx.send(Event::LoggedIn { greeting, own_ip });
        }
        server::ServerEvent::LoginFailed { reason } => {
            if let Some(tx) = inner.pending_login.lock().unwrap().take() {
                let _ = tx.send(Err(reason.clone()));
            }
            let _ = inner.event_tx.send(Event::LoginFailed { reason });
        }
        server::ServerEvent::UserAddress { username, addr } => {
            if addr.ip().is_unspecified() || addr.port() == 0 {
                // Offline; let the waiters time out.
                inner.pending_addresses.lock().unwrap().remove(&username);
                return;
            }
            inner
                .address_cache
                .lock()
                .unwrap()
                .insert(username.clone(), addr);
            if let Some(waiters) = inner.pending_addresses.lock().unwrap().remove(&username) {
                for w in waiters {
                    let _ = w.send(addr);
                }
            }
        }
        server::ServerEvent::ConnectToPeer {
            username,
            conn_type,
            addr,
            token,
            ..
        } => {
            // A peer that can't reach us directly asked the server to relay:
            // we connect out and pierce with their token.
            let inner = inner.clone();
            tokio::spawn(async move {
                let result = match conn_type.as_str() {
                    peer::CONN_TYPE_PEER => {
                        pierce_message_connection(inner, addr, token, username.clone()).await
                    }
                    peer::CONN_TYPE_TRANSFER => {
                        pierce_transfer_connection(inner, addr, token, username.clone()).await
                    }
                    _ => Ok(()), // "D" (distributed) lands in M5.
                };
                if let Err(e) = result {
                    log::debug!("pierce ({conn_type}) to {username} at {addr} failed: {e}");
                }
            });
        }
        server::ServerEvent::ExcludedSearchPhrases(phrases) => {
            *inner.excluded_phrases.write().unwrap() = phrases.clone();
            let _ = inner.event_tx.send(Event::ExcludedSearchPhrases(phrases));
        }
        server::ServerEvent::NetInfo(parents) => {
            crate::distributed::on_net_info(inner.clone(), parents);
        }
        server::ServerEvent::ServerSearch {
            username,
            token,
            query,
        } => {
            crate::responder::respond_to_search(inner.clone(), username, token, query);
        }
        server::ServerEvent::Embedded(bytes) => {
            // Server-wrapped distributed search.
            if let Ok(soulseek_proto::dist::DistributedEvent::SearchRequest {
                username,
                token,
                query,
            }) = soulseek_proto::dist::decode(bytes)
            {
                crate::responder::respond_to_search(inner.clone(), username, token, query);
            }
        }
        server::ServerEvent::CannotConnect { .. } | server::ServerEvent::Unhandled { .. } => {}
    }
}

// ---------------------------------------------------------------- listener

pub(crate) fn spawn_listener(
    listener: TcpListener,
    inner: Arc<Inner>,
    mut shutdown: mpsc::UnboundedReceiver<()>,
) {
    tokio::spawn(async move {
        loop {
            tokio::select! {
                // On shutdown, break so `listener` drops and its port is freed
                // — otherwise every reconnect leaks a bound port (WSAEADDRINUSE).
                _ = shutdown.recv() => {
                    log::debug!("listener shutting down, releasing port");
                    break;
                }
                res = listener.accept() => match res {
                    Ok((stream, addr)) => {
                        let inner = inner.clone();
                        tokio::spawn(async move {
                            if let Err(e) = handle_incoming(stream, inner).await {
                                log::debug!("peer {addr} dropped: {e}");
                            }
                        });
                    }
                    Err(e) => {
                        log::warn!("listener accept failed: {e}");
                        tokio::time::sleep(Duration::from_secs(1)).await;
                    }
                }
            }
        }
    });
}

/// Incoming connection: first frame identifies the peer and connection type.
async fn handle_incoming(mut stream: TcpStream, inner: Arc<Inner>) -> std::io::Result<()> {
    stream.set_nodelay(true).ok();
    let first = tokio::time::timeout(PEER_READ_TIMEOUT, read_frame(&mut stream))
        .await
        .map_err(|_| std::io::Error::other("init timeout"))??;
    let Some(first) = first else { return Ok(()) };

    match peer::decode_init(first) {
        Ok(peer::InitMessage::PeerInit {
            conn_type,
            username,
            ..
        }) => {
            match conn_type.as_str() {
                peer::CONN_TYPE_PEER => run_peer_connection(stream, inner, username).await,
                peer::CONN_TYPE_TRANSFER => {
                    // Inbound transfer connection: the peer identifies the
                    // transfer with 4 raw token bytes next.
                    let mut tok = [0u8; 4];
                    tokio::time::timeout(PEER_CONNECT_TIMEOUT, stream.read_exact(&mut tok))
                        .await
                        .map_err(|_| std::io::Error::other("token timeout"))??;
                    let token = u32::from_le_bytes(tok);
                    deliver_transfer_connection(&inner, &username, token, stream);
                    Ok(())
                }
                peer::CONN_TYPE_DISTRIBUTED => {
                    // A parent connected to us directly; run the child loop.
                    run_incoming_distributed(stream, inner).await
                }
                _ => Ok(()),
            }
        }
        Ok(peer::InitMessage::PierceFirewall { token }) => {
            // Answer to one of our indirect solicitations (P or F).
            let waiter = inner.pending_pierce.lock().unwrap().remove(&token);
            if let Some(tx) = waiter {
                let _ = tx.send(stream);
            }
            Ok(())
        }
        Err(_) => Ok(()),
    }
}

/// A distributed parent connected to us directly. Read its relayed
/// searches and answer them (same logic as an outbound parent).
async fn run_incoming_distributed(mut stream: TcpStream, inner: Arc<Inner>) -> std::io::Result<()> {
    inner
        .has_parent
        .store(true, std::sync::atomic::Ordering::Relaxed);
    let result = async {
        loop {
            let frame = match tokio::time::timeout(PEER_READ_TIMEOUT, read_frame(&mut stream)).await
            {
                Ok(f) => f?,
                Err(_) => return Ok(()),
            };
            let Some(frame) = frame else { return Ok(()) };
            if let Ok(soulseek_proto::dist::DistributedEvent::SearchRequest {
                username,
                token,
                query,
            }) = soulseek_proto::dist::decode(frame)
            {
                crate::responder::respond_to_search(inner.clone(), username, token, query);
            }
        }
    }
    .await;
    inner
        .has_parent
        .store(false, std::sync::atomic::Ordering::Relaxed);
    result
}

/// Hand an established inbound transfer connection to the download that
/// expects it.
fn deliver_transfer_connection(inner: &Arc<Inner>, username: &str, token: u32, stream: TcpStream) {
    let expected = inner
        .expected_transfers
        .lock()
        .unwrap()
        .remove(&(username.to_string(), token));
    match expected {
        Some(tx) => {
            let _ = tx.send(stream);
        }
        None => {
            log::debug!("unexpected transfer connection from {username} token {token}");
        }
    }
}

// ------------------------------------------------- indirect (pierce) roles

/// Server told us a peer wants a message connection: connect + pierce, then
/// treat like any peer message connection.
async fn pierce_message_connection(
    inner: Arc<Inner>,
    addr: SocketAddrV4,
    token: u32,
    username: String,
) -> std::io::Result<()> {
    let mut stream = tokio::time::timeout(PEER_CONNECT_TIMEOUT, TcpStream::connect(addr))
        .await
        .map_err(|_| std::io::Error::other("connect timeout"))??;
    stream.set_nodelay(true).ok();
    stream.write_all(&peer::pierce_firewall(token)).await?;
    run_peer_connection(stream, inner, username).await
}

/// Server told us a peer wants a transfer connection with us: connect,
/// pierce, read the 4-byte remote token, hand off to the waiting download.
async fn pierce_transfer_connection(
    inner: Arc<Inner>,
    addr: SocketAddrV4,
    token: u32,
    username: String,
) -> std::io::Result<()> {
    let mut stream = tokio::time::timeout(PEER_CONNECT_TIMEOUT, TcpStream::connect(addr))
        .await
        .map_err(|_| std::io::Error::other("connect timeout"))??;
    stream.set_nodelay(true).ok();
    stream.write_all(&peer::pierce_firewall(token)).await?;

    let mut tok = [0u8; 4];
    tokio::time::timeout(PEER_CONNECT_TIMEOUT, stream.read_exact(&mut tok))
        .await
        .map_err(|_| std::io::Error::other("token timeout"))??;
    let remote_token = u32::from_le_bytes(tok);
    deliver_transfer_connection(&inner, &username, remote_token, stream);
    Ok(())
}

// -------------------------------------------------- peer message connection

/// Run a peer message connection: writer drains a queue, reader routes
/// messages. Registers in the peers map for reuse; deregisters on exit.
async fn run_peer_connection(
    stream: TcpStream,
    inner: Arc<Inner>,
    username: String,
) -> std::io::Result<()> {
    let (mut rd, mut wr) = stream.into_split();
    let (write_tx, mut write_rx) = mpsc::unbounded_channel::<Vec<u8>>();

    inner.peers.lock().unwrap().insert(
        username.clone(),
        PeerHandle {
            write_tx: write_tx.clone(),
        },
    );

    let writer = tokio::spawn(async move {
        while let Some(frame) = write_rx.recv().await {
            if wr.write_all(&frame).await.is_err() {
                break;
            }
        }
    });

    let result = peer_message_loop(&mut rd, &inner, &username).await;

    // Only deregister if the map still points at this connection.
    {
        let mut peers = inner.peers.lock().unwrap();
        if let Some(handle) = peers.get(&username) {
            if handle.write_tx.same_channel(&write_tx) {
                peers.remove(&username);
            }
        }
    }
    writer.abort();
    result
}

async fn peer_message_loop(
    rd: &mut tokio::net::tcp::OwnedReadHalf,
    inner: &Arc<Inner>,
    username: &str,
) -> std::io::Result<()> {
    loop {
        let frame = match tokio::time::timeout(PEER_READ_TIMEOUT, read_frame(rd)).await {
            Ok(f) => f?,
            Err(_) => return Ok(()), // idle
        };
        let Some(mut body) = frame else { return Ok(()) };

        use bytes::Buf;
        if body.len() < 4 {
            continue;
        }
        let code = body.get_u32_le();

        match peer::decode(code, body) {
            Ok(ev) => route_peer_event(inner, username, ev),
            Err(e) => log::debug!("bad peer message {code} from {username}: {e}"),
        }
    }
}

fn route_peer_event(inner: &Arc<Inner>, username: &str, ev: peer::PeerEvent) {
    use peer::PeerEvent::*;
    match ev {
        SearchResponse(resp) => {
            let wanted = inner.active_searches.lock().unwrap().contains(&resp.token);
            if wanted && !resp.files.is_empty() {
                let _ = inner.event_tx.send(Event::SearchResult {
                    token: resp.token,
                    response: resp,
                });
            }
        }
        TransferResponse {
            token,
            allowed,
            size,
            message,
        } => {
            let waiter = inner
                .transfer_waits
                .lock()
                .unwrap()
                .remove(&(username.to_string(), token));
            if let Some(tx) = waiter {
                let _ = tx.send(TransferAck {
                    allowed,
                    size,
                    message,
                });
            }
        }
        TransferRequest {
            direction,
            token,
            filename,
            size,
        } => {
            if direction == peer::DIRECTION_UPLOAD {
                let waiter = inner
                    .upload_offers
                    .lock()
                    .unwrap()
                    .remove(&(username.to_string(), filename.clone()));
                match waiter {
                    Some(tx) => {
                        let _ = tx.send(Ok((token, size)));
                    }
                    None => log::debug!("unsolicited upload offer from {username}: {filename}"),
                }
            } else {
                // Their request to download one of our files → we upload.
                crate::upload::on_transfer_request(
                    inner.clone(),
                    username.to_string(),
                    token,
                    filename,
                );
            }
        }
        QueueDownload { filename } => {
            crate::upload::on_queue_download(inner.clone(), username.to_string(), filename);
        }
        PlaceInQueueRequest { filename } => {
            crate::upload::on_place_query(inner, username, &filename);
        }
        BrowseRequest => {
            if let Some(h) = inner.peers.lock().unwrap().get(username) {
                crate::responder::respond_browse(inner, h);
            }
        }
        UserInfoRequest => {
            if let Some(h) = inner.peers.lock().unwrap().get(username) {
                crate::responder::respond_user_info(inner, h);
            }
        }
        FolderContentsRequest { token, folder } => {
            if let Some(h) = inner.peers.lock().unwrap().get(username) {
                crate::responder::respond_folder_contents(inner, h, token, &folder);
            }
        }
        UploadDenied { filename, reason } => {
            let waiter = inner
                .upload_offers
                .lock()
                .unwrap()
                .remove(&(username.to_string(), filename));
            if let Some(tx) = waiter {
                let _ = tx.send(Err(reason));
            }
        }
        UploadFailed { filename } => {
            let waiter = inner
                .upload_offers
                .lock()
                .unwrap()
                .remove(&(username.to_string(), filename));
            if let Some(tx) = waiter {
                let _ = tx.send(Err("remote upload failed".into()));
            }
        }
        PlaceInQueueResponse { filename, place } => {
            crate::transfer::update_queue_place(inner, username, &filename, place);
        }
        Unhandled { code } => log::trace!("unhandled peer message {code} from {username}"),
    }
}

// ----------------------------------------------- outbound peer connections

/// Resolve a user's endpoint via the server (with cache).
pub(crate) async fn get_peer_address(inner: &Arc<Inner>, username: &str) -> Result<SocketAddrV4> {
    if let Some(addr) = inner.address_cache.lock().unwrap().get(username) {
        return Ok(*addr);
    }
    let (tx, rx) = oneshot::channel();
    inner
        .pending_addresses
        .lock()
        .unwrap()
        .entry(username.to_string())
        .or_default()
        .push(tx);
    inner.send_server(server::get_peer_address(username))?;

    tokio::time::timeout(Duration::from_secs(10), rx)
        .await
        .map_err(|_| ClientError::UserOffline(username.to_string()))?
        .map_err(|_| ClientError::UserOffline(username.to_string()))
}

/// Get (or establish) a message connection to a user. Tries direct first,
/// then falls back to a server-relayed indirect connection.
pub(crate) async fn get_or_connect_peer(inner: &Arc<Inner>, username: &str) -> Result<PeerHandle> {
    if let Some(h) = inner.peers.lock().unwrap().get(username) {
        return Ok(h.clone());
    }

    let addr = get_peer_address(inner, username).await?;

    // Direct attempt.
    let direct = async {
        let mut stream = tokio::time::timeout(PEER_CONNECT_TIMEOUT, TcpStream::connect(addr))
            .await
            .map_err(|_| std::io::Error::other("connect timeout"))??;
        stream.set_nodelay(true).ok();
        let init = peer::peer_init(
            &inner.own_username(),
            peer::CONN_TYPE_PEER,
            inner.next_token(),
        );
        stream.write_all(&init).await?;
        Ok::<TcpStream, std::io::Error>(stream)
    };

    let stream = match direct.await {
        Ok(s) => s,
        Err(e) => {
            log::debug!("direct connect to {username} failed ({e}); trying indirect");
            indirect_connect(inner, username, peer::CONN_TYPE_PEER).await?
        }
    };

    let inner2 = inner.clone();
    let username2 = username.to_string();
    tokio::spawn(async move {
        let _ = run_peer_connection(stream, inner2, username2).await;
    });

    // The handle is registered synchronously at the top of
    // run_peer_connection, but the spawn may not have run yet — poll briefly.
    for _ in 0..50 {
        if let Some(h) = inner.peers.lock().unwrap().get(username) {
            return Ok(h.clone());
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
    Err(ClientError::PeerUnreachable(username.to_string()))
}

/// Ask the server to have `username` connect to us; await the pierce.
pub(crate) async fn indirect_connect(
    inner: &Arc<Inner>,
    username: &str,
    conn_type: &str,
) -> Result<TcpStream> {
    let token = inner.next_token();
    let (tx, rx) = oneshot::channel();
    inner.pending_pierce.lock().unwrap().insert(token, tx);
    inner.send_server(server::connect_to_peer(token, username, conn_type))?;

    let out = tokio::time::timeout(INDIRECT_TIMEOUT, rx).await;
    inner.pending_pierce.lock().unwrap().remove(&token);
    out.map_err(|_| ClientError::PeerUnreachable(username.to_string()))?
        .map_err(|_| ClientError::PeerUnreachable(username.to_string()))
}

/// Establish an outbound transfer ("F") connection and negotiate the token.
pub(crate) async fn connect_transfer(
    inner: &Arc<Inner>,
    username: &str,
    token: u32,
) -> Result<TcpStream> {
    let addr = get_peer_address(inner, username).await?;

    let direct = async {
        let mut stream = tokio::time::timeout(PEER_CONNECT_TIMEOUT, TcpStream::connect(addr))
            .await
            .map_err(|_| std::io::Error::other("connect timeout"))??;
        stream.set_nodelay(true).ok();
        let init = peer::peer_init(&inner.own_username(), peer::CONN_TYPE_TRANSFER, token);
        stream.write_all(&init).await?;
        Ok::<TcpStream, std::io::Error>(stream)
    };

    let mut stream = match direct.await {
        Ok(s) => s,
        Err(e) => {
            log::debug!("direct F connect to {username} failed ({e}); trying indirect");
            indirect_connect(inner, username, peer::CONN_TYPE_TRANSFER).await?
        }
    };

    // Both flavors: identify the transfer with 4 raw token bytes.
    stream
        .write_all(&token.to_le_bytes())
        .await
        .map_err(|_| ClientError::PeerUnreachable(username.to_string()))?;
    Ok(stream)
}