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

//! Download state machine, mirroring Soulseek.NET's `DownloadToStreamAsync`:
//!
//! 1. Send `TransferRequest(download, token, filename)` on the P connection.
//! 2. Peer answers `TransferResponse`:
//!    - allowed → we initiate the F connection with our token.
//!    - "Queued" → wait for the peer's `TransferRequest(upload)`, ack it,
//!      await their inbound F connection (second-chance: initiate ours).
//! 3. On the F connection: write 8-byte start offset, then read raw bytes.

use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;

use soulseek_proto::peer;

use crate::{conn, ClientError, Event, Inner, Result};

/// How long we stay in the remote queue before giving up.
const QUEUE_TIMEOUT: Duration = Duration::from_secs(2 * 60 * 60);
/// Ack / connection-establishment timeout.
const NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(20);
/// No bytes for this long mid-transfer → fail.
const TRANSFER_STALL_TIMEOUT: Duration = Duration::from_secs(90);
/// Poll the remote queue position at this interval while queued.
const PLACE_POLL_INTERVAL: Duration = Duration::from_secs(60);

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DownloadState {
    Connecting,
    Requested,
    Queued { place: Option<u32> },
    Transferring,
    Complete { local_path: String },
    Failed { reason: String },
    Cancelled,
}

#[derive(Debug, Clone)]
pub struct DownloadUpdate {
    pub id: u64,
    pub username: String,
    /// Full remote path (backslash separated).
    pub filename: String,
    pub state: DownloadState,
    pub received: u64,
    pub total: u64,
    pub speed_bps: u64,
}

pub(crate) fn start_download(
    inner: Arc<Inner>,
    username: &str,
    filename: &str,
    size: u64,
    target: Option<PathBuf>,
) -> u64 {
    let id = inner
        .download_counter
        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    let cancel = CancellationToken::new();
    inner.downloads_by_file.lock().unwrap().insert(
        (username.to_string(), filename.to_string()),
        (id, cancel.clone()),
    );

    let username = username.to_string();
    let filename = filename.to_string();
    tokio::spawn(async move {
        let mut dl = Download {
            inner: inner.clone(),
            id,
            username: username.clone(),
            filename: filename.clone(),
            target,
            total: size,
            received: 0,
            speed_bps: 0,
            cancel: cancel.clone(),
        };

        let outcome = tokio::select! {
            r = dl.run() => r,
            _ = cancel.cancelled() => Err(ClientError::Timeout), // mapped below
        };

        match outcome {
            Ok(local_path) => dl.emit(DownloadState::Complete { local_path }),
            Err(_) if cancel.is_cancelled() => dl.emit(DownloadState::Cancelled),
            Err(e) => dl.emit(DownloadState::Failed {
                reason: e.to_string(),
            }),
        }

        inner
            .downloads_by_file
            .lock()
            .unwrap()
            .remove(&(username, filename));
    });
    id
}

/// Emit a queue-place update for a download (from PlaceInQueueResponse).
pub(crate) fn update_queue_place(inner: &Arc<Inner>, username: &str, filename: &str, place: u32) {
    let id = inner
        .downloads_by_file
        .lock()
        .unwrap()
        .get(&(username.to_string(), filename.to_string()))
        .map(|(id, _)| *id);
    if let Some(id) = id {
        let _ = inner.event_tx.send(Event::Download(DownloadUpdate {
            id,
            username: username.to_string(),
            filename: filename.to_string(),
            state: DownloadState::Queued { place: Some(place) },
            received: 0,
            total: 0,
            speed_bps: 0,
        }));
    }
}

struct Download {
    inner: Arc<Inner>,
    id: u64,
    username: String,
    filename: String,
    /// Explicit local target (stream-to-cache); None → downloads dir layout.
    target: Option<PathBuf>,
    total: u64,
    received: u64,
    speed_bps: u64,
    cancel: CancellationToken,
}

impl Download {
    fn emit(&self, state: DownloadState) {
        let _ = self.inner.event_tx.send(Event::Download(DownloadUpdate {
            id: self.id,
            username: self.username.clone(),
            filename: self.filename.clone(),
            state,
            received: self.received,
            total: self.total,
            speed_bps: self.speed_bps,
        }));
    }

    async fn run(&mut self) -> Result<String> {
        self.emit(DownloadState::Connecting);
        let handle = conn::get_or_connect_peer(&self.inner, &self.username).await?;

        // Register both waits before sending the request to avoid races:
        // the immediate ack, and the eventual upload offer if we get queued.
        let token = self.inner.next_token();
        let (ack_tx, ack_rx) = oneshot::channel();
        self.inner
            .transfer_waits
            .lock()
            .unwrap()
            .insert((self.username.clone(), token), ack_tx);

        let (offer_tx, offer_rx) = oneshot::channel();
        self.inner
            .upload_offers
            .lock()
            .unwrap()
            .insert((self.username.clone(), self.filename.clone()), offer_tx);

        handle
            .write_tx
            .send(peer::transfer_request_download(token, &self.filename))
            .map_err(|_| ClientError::PeerUnreachable(self.username.clone()))?;
        self.emit(DownloadState::Requested);

        let ack = tokio::time::timeout(NEGOTIATION_TIMEOUT, ack_rx)
            .await
            .map_err(|_| ClientError::Timeout)?
            .map_err(|_| ClientError::PeerUnreachable(self.username.clone()))?;

        let stream = if ack.allowed {
            // Free slot: we initiate the transfer connection with our token.
            if ack.size > 0 {
                self.total = ack.size;
            }
            self.emit(DownloadState::Transferring);
            conn::connect_transfer(&self.inner, &self.username, token).await?
        } else if ack
            .message
            .trim_end_matches('.')
            .eq_ignore_ascii_case("queued")
        {
            self.emit(DownloadState::Queued { place: None });
            self.await_queued_transfer(&handle, offer_rx).await?
        } else {
            return Err(ClientError::TransferRejected(ack.message));
        };

        self.receive(stream).await
    }

    /// We're in the remote queue: poll our place, wait for the peer's
    /// upload offer, ack it, and obtain the transfer connection.
    async fn await_queued_transfer(
        &mut self,
        handle: &crate::PeerHandle,
        offer_rx: oneshot::Receiver<crate::UploadOffer>,
    ) -> Result<TcpStream> {
        // Queue-place polling until the offer arrives.
        let poll_handle = handle.write_tx.clone();
        let poll_filename = self.filename.clone();
        let poller = tokio::spawn(async move {
            loop {
                tokio::time::sleep(PLACE_POLL_INTERVAL).await;
                if poll_handle
                    .send(peer::place_in_queue_request(&poll_filename))
                    .is_err()
                {
                    break;
                }
            }
        });

        let offer = tokio::time::timeout(QUEUE_TIMEOUT, offer_rx).await;
        poller.abort();

        let (remote_token, size) = offer
            .map_err(|_| ClientError::Timeout)?
            .map_err(|_| ClientError::PeerUnreachable(self.username.clone()))?
            .map_err(ClientError::TransferRejected)?;

        if size > 0 {
            self.total = size;
        }

        // Expect the peer's inbound F connection carrying their token —
        // registered before we ack so the connection can't race us.
        let (conn_tx, conn_rx) = oneshot::channel();
        self.inner
            .expected_transfers
            .lock()
            .unwrap()
            .insert((self.username.clone(), remote_token), conn_tx);

        // Re-resolve the peer connection in case the old one idled out.
        let handle = conn::get_or_connect_peer(&self.inner, &self.username).await?;
        handle
            .write_tx
            .send(peer::transfer_response_allowed(remote_token, self.total))
            .map_err(|_| ClientError::PeerUnreachable(self.username.clone()))?;

        self.emit(DownloadState::Transferring);

        match tokio::time::timeout(NEGOTIATION_TIMEOUT, conn_rx).await {
            Ok(Ok(stream)) => Ok(stream),
            _ => {
                // Second chance (Nicotine+ et al.): initiate from our end.
                self.inner
                    .expected_transfers
                    .lock()
                    .unwrap()
                    .remove(&(self.username.clone(), remote_token));
                log::debug!(
                    "no inbound F connection from {}; initiating second-chance",
                    self.username
                );
                conn::connect_transfer(&self.inner, &self.username, remote_token).await
            }
        }
    }

    /// Drive the transfer connection: write start offset, stream to disk.
    async fn receive(&mut self, mut stream: TcpStream) -> Result<String> {
        let (final_path, partial_path) = self.local_paths()?;

        // Resume from an existing partial file when its length is sane.
        let offset = match tokio::fs::metadata(&partial_path).await {
            Ok(m) if self.total > 0 && m.len() < self.total => m.len(),
            _ => 0,
        };

        let mut file = tokio::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .open(&partial_path)
            .await?;
        file.set_len(offset).await?;
        file.seek(std::io::SeekFrom::Start(offset)).await?;

        stream.write_all(&offset.to_le_bytes()).await?;
        self.received = offset;
        self.emit(DownloadState::Transferring);

        let mut buf = vec![0u8; 64 * 1024];
        let mut window_start = Instant::now();
        let mut window_bytes: u64 = 0;
        let mut last_emit = Instant::now();

        while self.total == 0 || self.received < self.total {
            let n = tokio::select! {
                _ = self.cancel.cancelled() => return Err(ClientError::Timeout),
                r = tokio::time::timeout(TRANSFER_STALL_TIMEOUT, stream.read(&mut buf)) => {
                    r.map_err(|_| ClientError::Timeout)??
                }
            };
            if n == 0 {
                if self.total == 0 {
                    break; // size unknown: EOF is completion
                }
                return Err(ClientError::PeerUnreachable(format!(
                    "{} closed the connection early ({}/{} bytes)",
                    self.username, self.received, self.total
                )));
            }
            file.write_all(&buf[..n]).await?;
            self.received += n as u64;
            window_bytes += n as u64;

            let elapsed = window_start.elapsed();
            if elapsed >= Duration::from_secs(1) {
                self.speed_bps = (window_bytes as f64 / elapsed.as_secs_f64()) as u64;
                window_start = Instant::now();
                window_bytes = 0;
            }
            if last_emit.elapsed() >= Duration::from_millis(250) {
                last_emit = Instant::now();
                self.emit(DownloadState::Transferring);
            }
        }

        file.flush().await?;
        drop(file);

        let final_path = unique_path(final_path).await;
        tokio::fs::rename(&partial_path, &final_path).await?;
        Ok(final_path.to_string_lossy().into_owned())
    }

    /// `<downloads_dir>/<user>/<remote folder leaf>/<file>` (+ `.partial`),
    /// or the explicit target when set.
    fn local_paths(&self) -> Result<(PathBuf, PathBuf)> {
        if let Some(target) = &self.target {
            if let Some(dir) = target.parent() {
                std::fs::create_dir_all(dir)?;
            }
            let partial = target.with_extension(format!(
                "{}partial",
                target
                    .extension()
                    .map(|e| format!("{}.", e.to_string_lossy()))
                    .unwrap_or_default()
            ));
            return Ok((target.clone(), partial));
        }
        let (folder, name) = match self.filename.rsplit_once('\\') {
            Some((dir, n)) => (dir.rsplit('\\').next().unwrap_or(""), n),
            None => ("", self.filename.as_str()),
        };
        let mut dir = self.inner.downloads_dir.clone();
        dir.push(sanitize(&self.username));
        if !folder.is_empty() {
            dir.push(sanitize(folder));
        }
        std::fs::create_dir_all(&dir)?;
        let final_path = dir.join(sanitize(name));
        let partial = dir.join(format!("{}.partial", sanitize(name)));
        Ok((final_path, partial))
    }
}

/// Make a remote path component safe for Windows.
fn sanitize(part: &str) -> String {
    let cleaned: String = part
        .chars()
        .map(|c| match c {
            '<' | '>' | ':' | '"' | '/' | '\\' | '|' | '?' | '*' => '_',
            c if (c as u32) < 0x20 => '_',
            c => c,
        })
        .collect();
    cleaned.trim_end_matches([' ', '.']).to_string()
}

/// Avoid clobbering an existing file: append " (2)", " (3)", …
async fn unique_path(path: PathBuf) -> PathBuf {
    if tokio::fs::metadata(&path).await.is_err() {
        return path;
    }
    let stem = path
        .file_stem()
        .map(|s| s.to_string_lossy().into_owned())
        .unwrap_or_default();
    let ext = path
        .extension()
        .map(|e| format!(".{}", e.to_string_lossy()))
        .unwrap_or_default();
    let dir = path.parent().map(PathBuf::from).unwrap_or_default();
    for i in 2.. {
        let candidate = dir.join(format!("{stem} ({i}){ext}"));
        if tokio::fs::metadata(&candidate).await.is_err() {
            return candidate;
        }
    }
    unreachable!()
}

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

    #[test]
    fn sanitize_windows_reserved() {
        assert_eq!(sanitize("AC/DC: Back in Black?"), "AC_DC_ Back in Black_");
        assert_eq!(sanitize("trailing dots..."), "trailing dots");
    }
}