Merge pull request #103 from michel/michel/fix/brokered-upload-connection

fix(lib): serve uploads to downloaders we cannot dial

soulseek-rs-lib/src/client/context.rs

@@ -5,7 +5,7 @@
 //! of them having to be widened for the move.
 
 use super::{
-    ActorSystem, Arc, AtomicUsize, BROKER_CONNECT_TIMEOUT,
+    ActorSystem, Arc, AtomicBool, AtomicUsize, BROKER_CONNECT_TIMEOUT,
     BROWSE_PROTECT_WINDOW, ClientContext, DEFAULT_MAX_PEERS,
     DEFAULT_UPLOAD_SLOTS, DownloadStatus, DownloadStore, HashMap, HashSet,
     Instant, PENDING_PEER_TTL, Recommendation, RoomEvent, RoomInfo, RoomTicker,
@@ -50,6 +50,7 @@ impl ClientContext {
             uploads: HashMap::new(),
             active_uploads: HashMap::new(),
             pending_serves: HashMap::new(),
+            brokered_uploads: HashMap::new(),
             browse_results: HashMap::new(),
             folder_contents: HashMap::new(),
             peer_infos: HashMap::new(),
@@ -707,6 +708,45 @@ impl ClientContext {
             .map(|(username, _)| username)
     }
 
+    /// Returns the claim flag the dialling route races against.
+    pub(crate) fn add_brokered_upload(
+        &mut self,
+        connect_token: u32,
+        token: u32,
+        real_path: std::path::PathBuf,
+    ) -> Arc<AtomicBool> {
+        let claimed = Arc::new(AtomicBool::new(false));
+        self.brokered_uploads.insert(
+            connect_token,
+            super::BrokeredUpload {
+                token,
+                real_path,
+                deadline: Instant::now() + BROKER_CONNECT_TIMEOUT,
+                claimed: claimed.clone(),
+            },
+        );
+        claimed
+    }
+
+    pub(crate) fn take_brokered_upload(
+        &mut self,
+        connect_token: u32,
+    ) -> Option<super::BrokeredUpload> {
+        self.brokered_uploads.remove(&connect_token)
+    }
+
+    /// Uploads whose downloader never connected back in time.
+    pub(crate) fn take_expired_brokered_uploads(
+        &mut self,
+        now: Instant,
+    ) -> Vec<super::BrokeredUpload> {
+        let (expired, live) = std::mem::take(&mut self.brokered_uploads)
+            .into_iter()
+            .partition(|(_, upload)| now >= upload.deadline);
+        self.brokered_uploads = live;
+        expired.into_values().collect()
+    }
+
     pub fn take_expired_connects(&mut self, now: Instant) -> Vec<String> {
         let mut expired = Vec::new();
         self.pending_connect_tokens

soulseek-rs-lib/src/client/mod.rs

@@ -114,6 +114,24 @@ struct UploadJob {
     offered: Option<Instant>,
 }
 
+/// An upload waiting for its downloader to connect back to our listener while
+/// we also dial them. Whichever socket arrives first claims the transfer.
+pub(crate) struct BrokeredUpload {
+    pub(crate) token: u32,
+    pub(crate) real_path: std::path::PathBuf,
+    deadline: Instant,
+    claimed: Arc<AtomicBool>,
+}
+
+impl BrokeredUpload {
+    /// Take the transfer for this route; false if another route already did.
+    pub(crate) fn claim(&self) -> bool {
+        self.claimed
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_ok()
+    }
+}
+
 /// Live bookkeeping for an upload being served (or recently finished).
 struct ActiveUpload {
     username: String,
@@ -504,6 +522,10 @@ pub struct ClientContext {
     active_uploads: HashMap<u32, ActiveUpload>,
     /// Upload tokens waiting for the downloader's address to be resolved.
     pending_serves: HashMap<String, Vec<u32>>,
+    /// Uploads whose downloader we could not dial, keyed by the correlation
+    /// token we sent in `ConnectToPeer(F)`; the peer quotes it back in a
+    /// `PierceFirewall` and we serve over that socket instead.
+    brokered_uploads: HashMap<u32, BrokeredUpload>,
     /// Shared-file listings received from peers we browsed.
     browse_results: HashMap<String, Vec<SharedDirectory>>,
     /// One-folder listings received from peers, keyed by peer and folder.

soulseek-rs-lib/src/client/operations/mod.rs

@@ -151,26 +151,26 @@ impl Client {
                         // its failure is exactly what makes an
                         // unreachable peer fall back to the
                         // server-brokered path.
-                        let waiting_serves = if port == 0 {
+                        if port == 0 {
                             warn!(
                                 "[client] server reports no listening port for {}",
                                 username
                             );
-                            Vec::new()
+                        }
-                        } else {
+                        let waiting_serves = match client_context.write_safe() {
-                            match client_context.write_safe() {
+                            Ok(mut ctx) => {
-                                Ok(mut ctx) => {
+                                if port != 0 {
                                     ctx.cache_peer_address(
                                         &username,
                                         host.clone(),
                                         port,
                                     );
-                                    ctx.pending_serves
-                                        .remove(&username)
-                                        .unwrap_or_default()
                                 }
-                                Err(_) => Vec::new(),
+                                ctx.pending_serves
+                                    .remove(&username)
+                                    .unwrap_or_default()
                             }
+                            Err(_) => Vec::new(),
                         };
                         for token in waiting_serves {
                             Self::spawn_serve(
@@ -622,16 +622,29 @@ impl Client {
                         // token is the correlation we sent, so it names the
                         // peer: end the wait now instead of letting every
                         // queued download for them sit until the timeout.
-                        let username = match client_context.write_safe() {
+                        let (username, upload) =
-                            Ok(mut ctx) => ctx.take_pending_connect(token),
+                            match client_context.write_safe() {
-                            Err(e) => {
+                                Ok(mut ctx) => (
-                                error!(
+                                    ctx.take_pending_connect(token),
-                                    "[client] CantConnectToPeer write: {}",
+                                    ctx.take_brokered_upload(token),
-                                    e
+                                ),
-                                );
+                                Err(e) => {
-                                None
+                                    error!(
-                            }
+                                        "[client] CantConnectToPeer write: {}",
-                        };
+                                        e
+                                    );
+                                    (None, None)
+                                }
+                            };
+                        if let Some(upload) = upload
+                            && upload.claim()
+                        {
+                            Self::fail_upload(
+                                &client_context,
+                                upload.token,
+                                "peer could not connect back",
+                            );
+                        }
                         if let Some(username) = username {
                             debug!(
                                 "[client] {} cannot connect back (token {})",
@@ -788,5 +801,16 @@ impl Client {
         for username in expired {
             Self::fail_queued_downloads(client_context, &username);
         }
+        let expired = match client_context.write_safe() {
+            Ok(mut ctx) => ctx.take_expired_brokered_uploads(Instant::now()),
+            Err(_) => return,
+        };
+        for upload in expired.iter().filter(|u| u.claim()) {
+            Self::fail_upload(
+                client_context,
+                upload.token,
+                "peer never connected back",
+            );
+        }
     }
 }

soulseek-rs-lib/src/client/uploads.rs

@@ -1,11 +1,14 @@
 use super::{
-    ActiveUpload, Arc, Client, ClientContext, DownloadStatus, PeerMessage,
+    ActiveUpload, Arc, BrokeredUpload, Client, ClientContext, DownloadStatus,
-    RwLock, RwLockExt, ServerMessage, collect_failed_tokens, error,
+    PeerMessage, RwLock, RwLockExt, ServerMessage, collect_failed_tokens,
-    next_upload_token, thread,
+    debug, error, next_connect_token, next_upload_token, thread,
 };
 use crate::message::server::MessageFactory;
+use crate::peer::{ConnectionType, upload_peer};
 use crate::types::UploadStatus;
-use std::sync::atomic::{AtomicBool, AtomicU64};
+use std::net::TcpStream;
+use std::path::Path;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 
 impl Client {
     /// Whether the server listed `username` as privileged. Privileged peers are
@@ -136,55 +139,190 @@ impl Client {
         drop(ctx);
         let own = own_username.to_string();
         let real_path = job.real_path;
+        let downloader = job.downloader;
         let context = client_context.clone();
         thread::spawn(move || {
-            let result = crate::peer::upload_peer::serve_file(
+            // Most downloaders sit behind a NAT, so dialling them is a gamble
-                &host,
+            // that costs a 20s timeout when it fails. Like the reference
-                port,
+            // clients, ask the server to have them connect back at the same
+            // time; the first socket to arrive claims the transfer, and port 0
+            // (the server not knowing where they listen) skips the dial.
+            let Some((connect_token, claimed)) =
+                Self::broker_upload(&context, token, &real_path, &downloader)
+            else {
+                return;
+            };
+            if port == 0 {
+                return;
+            }
+            let stream = match upload_peer::dial(&host, port) {
+                Ok(stream) => stream,
+                Err(e) => {
+                    debug!(
+                        "[client] cannot dial {downloader} at {host}:{port} \
+                         for upload {token} ({e}); waiting for them to connect"
+                    );
+                    return;
+                }
+            };
+            if claimed
+                .compare_exchange(
+                    false,
+                    true,
+                    Ordering::AcqRel,
+                    Ordering::Acquire,
+                )
+                .is_err()
+            {
+                return;
+            }
+            if let Ok(mut ctx) = context.write_safe() {
+                ctx.take_brokered_upload(connect_token);
+            }
+            let result = upload_peer::serve_direct(
+                stream,
                 &own,
                 token,
                 &real_path,
                 &bytes_sent,
                 &cancel,
             );
-            let streamed = result.as_ref().ok().copied();
+            Self::finish_upload(&context, token, &real_path, result);
-            let status = match &result {
-                Ok(_) => UploadStatus::Completed,
-                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
-                    UploadStatus::Cancelled
-                }
-                Err(e) => {
-                    error!("[client] serve {}: {}", real_path.display(), e);
-                    UploadStatus::Failed(e.to_string())
-                }
-            };
-            let mut report = None;
-            if let Ok(mut ctx) = context.write_safe()
-                && let Some(upload) = ctx.active_uploads.get_mut(&token)
-            {
-                let secs = upload.started.elapsed().as_secs_f64();
-                upload.status = status;
-                if let Some(bytes) = streamed
-                    && secs > 0.0
-                {
-                    let speed = (bytes as f64 / secs) as u32;
-                    ctx.last_upload_speed = speed;
-                    report = ctx.server_sender.clone().map(|s| (s, speed));
-                }
-            }
-            // The server folds each finished upload's rate into the average
-            // speed it shows other users for us.
-            if let Some((server, speed)) = report {
-                let _ = server.send(ServerMessage::SendMessage(
-                    MessageFactory::build_send_upload_speed(speed),
-                ));
-            }
-            // The slot this transfer held is free now, so whoever is next in
-            // line gets it without waiting for another request to arrive.
-            Self::pump_upload_queue(&context);
         });
     }
 
+    /// Park the upload under a correlation token and ask the server to have
+    /// `downloader` connect back to us with it. Returns the token and the claim
+    /// flag the dialling route races against; `None` fails the upload.
+    fn broker_upload(
+        context: &Arc<RwLock<ClientContext>>,
+        token: u32,
+        real_path: &Path,
+        downloader: &str,
+    ) -> Option<(u32, Arc<AtomicBool>)> {
+        let connect_token = next_connect_token();
+        let (claimed, sender) = match context.write_safe() {
+            Ok(mut ctx) => (
+                ctx.add_brokered_upload(
+                    connect_token,
+                    token,
+                    real_path.to_path_buf(),
+                ),
+                ctx.server_sender.clone(),
+            ),
+            Err(_) => return None,
+        };
+        let Some(sender) = sender else {
+            Self::fail_upload(context, token, "not connected to the server");
+            return None;
+        };
+        let _ = sender.send(ServerMessage::SendMessage(
+            MessageFactory::build_connect_to_peer(
+                connect_token,
+                downloader,
+                ConnectionType::F,
+            ),
+        ));
+        Some((connect_token, claimed))
+    }
+
+    /// The downloader of a brokered upload connected back: stream the file
+    /// over its socket on a background thread, unless the dial beat it.
+    pub(crate) fn serve_pierced(
+        client_context: &Arc<RwLock<ClientContext>>,
+        upload: BrokeredUpload,
+        stream: TcpStream,
+    ) {
+        if !upload.claim() {
+            return;
+        }
+        let Some((bytes_sent, cancel)) =
+            client_context.write_safe().ok().and_then(|mut ctx| {
+                let active = ctx.active_uploads.get_mut(&upload.token)?;
+                active.started = std::time::Instant::now();
+                Some((active.bytes_sent.clone(), active.cancel.clone()))
+            })
+        else {
+            return;
+        };
+        let context = client_context.clone();
+        thread::spawn(move || {
+            let result = upload_peer::serve_pierced(
+                stream,
+                upload.token,
+                &upload.real_path,
+                &bytes_sent,
+                &cancel,
+            );
+            Self::finish_upload(
+                &context,
+                upload.token,
+                &upload.real_path,
+                result,
+            );
+        });
+    }
+
+    /// Mark upload `token` failed and hand its slot to the next in line.
+    pub(crate) fn fail_upload(
+        client_context: &Arc<RwLock<ClientContext>>,
+        token: u32,
+        reason: &str,
+    ) {
+        if let Ok(mut ctx) = client_context.write_safe()
+            && let Some(upload) = ctx.active_uploads.get_mut(&token)
+        {
+            error!("[client] upload {token} to {}: {reason}", upload.username);
+            upload.status = UploadStatus::Failed(reason.to_string());
+        }
+        Self::pump_upload_queue(client_context);
+    }
+
+    /// Record how a served upload ended, report its rate to the server and free
+    /// its slot.
+    fn finish_upload(
+        context: &Arc<RwLock<ClientContext>>,
+        token: u32,
+        real_path: &Path,
+        result: std::io::Result<u64>,
+    ) {
+        let streamed = result.as_ref().ok().copied();
+        let status = match &result {
+            Ok(_) => UploadStatus::Completed,
+            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
+                UploadStatus::Cancelled
+            }
+            Err(e) => {
+                error!("[client] serve {}: {}", real_path.display(), e);
+                UploadStatus::Failed(e.to_string())
+            }
+        };
+        let mut report = None;
+        if let Ok(mut ctx) = context.write_safe()
+            && let Some(upload) = ctx.active_uploads.get_mut(&token)
+        {
+            let secs = upload.started.elapsed().as_secs_f64();
+            upload.status = status;
+            if let Some(bytes) = streamed
+                && secs > 0.0
+            {
+                let speed = (bytes as f64 / secs) as u32;
+                ctx.last_upload_speed = speed;
+                report = ctx.server_sender.clone().map(|s| (s, speed));
+            }
+        }
+        // The server folds each finished upload's rate into the average
+        // speed it shows other users for us.
+        if let Some((server, speed)) = report {
+            let _ = server.send(ServerMessage::SendMessage(
+                MessageFactory::build_send_upload_speed(speed),
+            ));
+        }
+        // The slot this transfer held is free now, so whoever is next in
+        // line gets it without waiting for another request to arrive.
+        Self::pump_upload_queue(context);
+    }
+
     pub(crate) fn process_failed_uploads(
         client_context: Arc<RwLock<ClientContext>>,
         username: &str,

soulseek-rs-lib/src/peer/listen.rs

@@ -261,7 +261,18 @@ fn handle_pierce_firewall(
     let token = message.read_int32();
 
     let username = match context.client_context.write_safe() {
-        Ok(mut ctx) => ctx.take_pending_connect(token),
+        Ok(mut ctx) => {
+            if let Some(upload) = ctx.take_brokered_upload(token) {
+                drop(ctx);
+                crate::client::Client::serve_pierced(
+                    &context.client_context,
+                    upload,
+                    stream,
+                );
+                return;
+            }
+            ctx.take_pending_connect(token)
+        }
         Err(e) => {
             error!("[listener] pierce firewall lock: {}", e);
             return;

soulseek-rs-lib/src/peer/upload_peer.rs

@@ -24,6 +24,7 @@ use crate::trace;
 ///
 /// # Errors
 /// Returns any I/O error opening the file or talking to the peer.
+#[cfg(test)]
 pub fn serve_file(
     host: &str,
     port: u32,
@@ -33,16 +34,28 @@ pub fn serve_file(
     bytes_sent: &AtomicU64,
     cancel: &AtomicBool,
 ) -> io::Result<u64> {
-    let mut file = File::open(path)?;
+    let stream = dial(host, port)?;
+    serve_direct(stream, own_username, token, path, bytes_sent, cancel)
+}
 
+/// Connect to `host:port` with the transfer deadlines applied.
+///
+/// # Errors
+/// Returns the connect error, or `InvalidInput` for an unresolvable address.
+pub fn dial(host: &str, port: u32) -> io::Result<TcpStream> {
     let socket = format!("{host}:{port}")
         .to_socket_addrs()?
         .next()
         .ok_or_else(|| {
             io::Error::new(io::ErrorKind::InvalidInput, "no address")
         })?;
-    let mut stream =
+    with_deadlines(TcpStream::connect_timeout(
-        TcpStream::connect_timeout(&socket, Duration::from_secs(20))?;
+        &socket,
+        Duration::from_secs(20),
+    )?)
+}
+
+fn with_deadlines(stream: TcpStream) -> io::Result<TcpStream> {
     // A downloader that connects and then goes silent must cost seconds, not
     // the rest of the session: without deadlines a wedged `read_exact` (or a
     // zero-window `write_all`) holds this thread, the socket, the open file,
@@ -55,7 +68,23 @@ pub fn serve_file(
     stream.set_read_timeout(Some(Duration::from_secs(30)))?;
     stream.set_write_timeout(Some(Duration::from_secs(30)))?;
     stream.set_nodelay(true).ok();
+    Ok(stream)
+}
 
+/// Stream `path` over a connection we dialled: announce with `PeerInit(F)`
+/// and the token, then hand over to [`stream_file`].
+///
+/// # Errors
+/// Returns any I/O error opening the file or talking to the peer.
+pub fn serve_direct(
+    mut stream: TcpStream,
+    own_username: &str,
+    token: u32,
+    path: &Path,
+    bytes_sent: &AtomicU64,
+    cancel: &AtomicBool,
+) -> io::Result<u64> {
+    let file = File::open(path)?;
     // PeerInit(F) + the 4-byte token in a single write so they coalesce.
     let mut init = MessageFactory::build_peer_init_message(
         own_username,
@@ -66,7 +95,36 @@ pub fn serve_file(
     init.extend_from_slice(&token.to_le_bytes());
     stream.write_all(&init)?;
     stream.flush()?;
+    stream_file(stream, file, path, bytes_sent, cancel)
+}
 
+/// Stream `path` over a connection the downloader opened to us after the
+/// server brokered it: it has already sent `PierceFirewall`, so only the
+/// 4-byte transfer token precedes the offset exchange.
+///
+/// # Errors
+/// Returns any I/O error opening the file or talking to the peer.
+pub fn serve_pierced(
+    stream: TcpStream,
+    token: u32,
+    path: &Path,
+    bytes_sent: &AtomicU64,
+    cancel: &AtomicBool,
+) -> io::Result<u64> {
+    let file = File::open(path)?;
+    let mut stream = with_deadlines(stream)?;
+    stream.write_all(&token.to_le_bytes())?;
+    stream.flush()?;
+    stream_file(stream, file, path, bytes_sent, cancel)
+}
+
+fn stream_file(
+    mut stream: TcpStream,
+    mut file: File,
+    path: &Path,
+    bytes_sent: &AtomicU64,
+    cancel: &AtomicBool,
+) -> io::Result<u64> {
     // The downloader replies with an 8-byte START_DOWNLOAD offset first. A
     // non-zero offset means it is resuming and already holds that prefix, so
     // seek past it and count it as sent to keep progress absolute.
@@ -104,7 +162,11 @@ pub fn serve_file(
     // and a peer that leaves the socket open, or whose FIN a NAT swallows,
     // would hold it for the whole wait. Nicotine+ finishes the transfer at
     // delivery for the same reason; the close below is ours to make.
-    trace!("[upload] served {} to {}:{}", path.display(), host, port);
+    trace!(
+        "[upload] served {} to {:?}",
+        path.display(),
+        stream.peer_addr()
+    );
     Ok(streamed)
 }
 

soulseek-rs-lib/tests/e2e.rs

@@ -482,8 +482,8 @@ fn registered_username_can_relogin_with_same_password() {
 // ---------------------------------------------------------------------------
 // Peer-to-peer download coverage.
 //
-// The client cannot serve files (the upload side is not implemented), so a real
+// To exercise the download path on its own, the peer is a minimal in-process
-// download is exercised with a minimal in-process "mock uploader" that speaks
+// "mock uploader" rather than a second client. It speaks
 // the peer protocol using the library's own public `Message` wire format. It
 // drives the same path a real peer would after a search: a `P` control
 // connection to our listener, a `QueueUpload` → `TransferRequest` →
@@ -5692,3 +5692,65 @@ fn a_cancel_ends_a_login_waiting_for_its_verdict_and_nothing_follows_it() {
         0
     );
 }
+
+#[test]
+fn a_firewalled_downloader_still_receives_the_upload() {
+    // Issue #102: the leecher cannot be dialled (no listener, as behind a NAT),
+    // so the sharer must ask the server to broker the F connection and serve
+    // the file over the leecher's PierceFirewall connection.
+    let server = server_or_skip!();
+
+    let share_dir = unique_download_dir();
+    let content: Vec<u8> = (0..4096u32).map(|i| (i % 251) as u8).collect();
+    let filename = "e2e_probe_fwleech.bin";
+    std::fs::write(share_dir.join(filename), &content).unwrap();
+
+    let mut sharer = Client::with_settings(ClientSettings {
+        shared_directories: vec![share_dir.display().to_string()],
+        ..server.listening_settings(
+            "e2e_fw_sharer",
+            "pw",
+            free_port().expect("sharer port"),
+        )
+    });
+    sharer.connect().expect("sharer connect");
+    assert!(sharer.login().expect("sharer login"));
+
+    let mut leecher =
+        Client::with_settings(server.settings("e2e_fw_leecher", "pw"));
+    leecher.connect().expect("leecher connect");
+    assert!(leecher.login().expect("leecher login"));
+    std::thread::sleep(Duration::from_secs(1));
+
+    let query = "fwleech";
+    let _ = leecher.search(query, Duration::from_secs(3));
+    let (result_path, size) = reply_from(&leecher, query, "e2e_fw_sharer")
+        .and_then(|reply| reply.files.into_iter().next())
+        .map(|file| (file.name, file.size))
+        .expect("leecher should find the sharer's file");
+
+    let download_dir = unique_download_dir();
+    let (_download, status_rx) = leecher
+        .download(
+            result_path.clone(),
+            "e2e_fw_sharer".to_string(),
+            size,
+            download_dir.display().to_string(),
+        )
+        .expect("start download");
+
+    let completed =
+        wait_for_completion(&leecher, &status_rx, Duration::from_secs(40));
+    assert!(
+        completed,
+        "the download to a firewalled leecher should complete"
+    );
+
+    let basename = result_path.rsplit(['\\', '/']).next().unwrap();
+    assert_eq!(
+        std::fs::read(download_dir.join(basename)).expect("downloaded file"),
+        content
+    );
+    let _ = std::fs::remove_dir_all(&download_dir);
+    let _ = std::fs::remove_dir_all(&share_dir);
+}