fix(peer): wait for the transfer token before the download lookup

The listener looked the download up immediately after the PeerInit(F),
but the 4-byte transfer token routinely arrives in a separate TCP
segment, so every firewalled peer download failed with "Download info
missing". Wait for the token within the handshake deadline and give up
on the connection when it never arrives.

Also quiets the dispatcher's unhandled-message-code log to debug; the
server routinely sends codes the library doesn't implement.

Fixes #41

soulseek-rs-lib/src/dispatcher.rs

@@ -1,7 +1,7 @@
 use crate::message::{Message, handlers::Handlers};
 use std::sync::mpsc::Sender;
 
-use crate::warn;
+use crate::debug;
 
 pub struct MessageDispatcher<Op> {
     owner_name: String,
@@ -30,7 +30,7 @@ impl<Op> MessageDispatcher<Op> {
             message.set_pointer(8);
             handler.handle(message, self.sender.clone());
         } else {
-            warn!(
+            debug!(
                 "[{}:dispatcher] No handler found for message code: {code}",
                 self.owner_name
             );

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

@@ -96,12 +96,32 @@ fn parse_peer_init_message(mut message: Message) -> Option<PeerInitData> {
     })
 }
 
+fn wait_for_transfer_token(
+    stream: &mut TcpStream,
+    reader: &mut MessageReader,
+    deadline: Instant,
+) -> io::Result<()> {
+    while reader.buffer_len() < 4 {
+        let remaining = deadline.saturating_duration_since(Instant::now());
+        if remaining.is_zero() {
+            return Err(io::Error::new(
+                io::ErrorKind::TimedOut,
+                "transfer token deadline passed",
+            ));
+        }
+        stream.set_read_timeout(Some(remaining))?;
+        reader.read_from_socket(stream)?;
+    }
+    stream.set_read_timeout(None)?;
+    Ok(())
+}
+
 fn extract_download_from_buffer(
     reader: &mut MessageReader,
     client_context: &Arc<RwLock<ClientContext>>,
     username: &str,
     peer_ip: &str,
-    peer_port: u16,
+    peer_port: u32,
 ) -> Option<Download> {
     if reader.buffer_len() == 0 {
         return None;
@@ -161,24 +181,32 @@ fn handle_peer_connection(
 
 fn handle_file_connection(
     peer: Peer,
-    stream: TcpStream,
+    mut stream: TcpStream,
     mut reader: MessageReader,
     token: u32,
     context: &ConnectionContext,
-    peer_ip: &str,
-    peer_port: u16,
+    deadline: Instant,
 ) {
     trace!(
         "[client] DownloadFromPeer token: {} peer: {:?}",
         token, peer
     );
 
+    if let Err(e) = wait_for_transfer_token(&mut stream, &mut reader, deadline)
+    {
+        debug!(
+            "[listener:{}:{}] no transfer token arrived: {e}",
+            peer.host, peer.port
+        );
+        return;
+    }
+
     let download = extract_download_from_buffer(
         &mut reader,
         &context.client_context,
         &peer.username,
-        peer_ip,
-        peer_port,
+        &peer.host,
+        peer.port,
     );
     let failure_token = download.as_ref().map(|d| d.token);
 
@@ -368,8 +396,7 @@ fn handle_incoming_connection(
             reader,
             init_data.token,
             &context,
-            &peer_ip,
-            peer_port,
+            deadline,
         ),
         ConnectionType::D => {
             debug!(
@@ -482,6 +509,60 @@ mod tests {
         let _ = writer.join();
     }
 
+    #[test]
+    fn the_transfer_token_is_waited_for_across_segments() {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let addr = listener.local_addr().unwrap();
+        let writer = std::thread::spawn(move || {
+            let mut socket = TcpStream::connect(addr).unwrap();
+            std::thread::sleep(Duration::from_millis(100));
+            socket.write_all(&42u32.to_le_bytes()).unwrap();
+            std::thread::sleep(Duration::from_millis(200));
+        });
+        let (mut stream, _) = listener.accept().unwrap();
+        let mut reader = MessageReader::new();
+
+        let result = wait_for_transfer_token(
+            &mut stream,
+            &mut reader,
+            Instant::now() + Duration::from_secs(5),
+        );
+
+        assert!(
+            result.is_ok(),
+            "token in a later segment must be waited for"
+        );
+        assert!(reader.buffer_len() >= 4, "token bytes must be buffered");
+        let _ = writer.join();
+    }
+
+    #[test]
+    fn waiting_for_a_transfer_token_stops_at_the_deadline() {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let addr = listener.local_addr().unwrap();
+        let writer = std::thread::spawn(move || {
+            let socket = TcpStream::connect(addr).unwrap();
+            std::thread::sleep(Duration::from_millis(500));
+            drop(socket);
+        });
+        let (mut stream, _) = listener.accept().unwrap();
+        let mut reader = MessageReader::new();
+
+        let started = Instant::now();
+        let result = wait_for_transfer_token(
+            &mut stream,
+            &mut reader,
+            started + Duration::from_millis(200),
+        );
+
+        assert!(result.is_err(), "a silent peer must not park the thread");
+        assert!(
+            started.elapsed() < Duration::from_secs(5),
+            "the deadline must bound the wait"
+        );
+        let _ = writer.join();
+    }
+
     #[test]
     fn an_oversized_handshake_is_rejected() {
         let listener = TcpListener::bind("127.0.0.1:0").unwrap();

soulseek-rs-lib/tests/e2e.rs

@@ -478,6 +478,7 @@ struct MockUpload {
     content: Vec<u8>,
     token: u32,
     ready: Sender<()>,
+    token_delay: Duration,
 }
 
 /// Build a `PeerInit` (peer code 1) frame: `[len][1][username][conn_type][token]`.
@@ -562,6 +563,7 @@ fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
         &cfg.peer_username,
         cfg.token,
         &cfg.content,
+        cfg.token_delay,
     )
 }
 
@@ -589,10 +591,8 @@ fn run_refusing_uploader(
 }
 
 /// Open an F (file transfer) connection to `downloader_addr` and stream
-/// `content`. PeerInit(F) and the 4-byte transfer token are written together so
-/// the token lands in the listener's read buffer, where the download is looked
-/// up by token; the downloader then sends an 8-byte START_DOWNLOAD offset before
-/// we send the bytes.
+/// `content`; the downloader sends an 8-byte START_DOWNLOAD offset before we
+/// send the bytes.
 ///
 /// Like a real peer, only the bytes past that offset are sent. Returns the
 /// offset the downloader asked for.
@@ -601,12 +601,19 @@ fn serve_file_over_f(
     username: &str,
     token: u32,
     content: &[u8],
+    token_delay: Duration,
 ) -> std::io::Result<u64> {
     let mut f = connect_retry(downloader_addr, Duration::from_secs(5))?;
     f.set_read_timeout(Some(Duration::from_secs(10)))?;
+    if token_delay.is_zero() {
         let mut init = peer_init_bytes(username, "F", token);
         init.extend_from_slice(&token.to_le_bytes());
         f.write_all(&init)?;
+    } else {
+        f.write_all(&peer_init_bytes(username, "F", token))?;
+        std::thread::sleep(token_delay);
+        f.write_all(&token.to_le_bytes())?;
+    }
     f.flush()?;
 
     let mut start = [0u8; 8];
@@ -717,14 +724,13 @@ fn unique_download_dir() -> PathBuf {
     dir
 }
 
-#[test]
-fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
+fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
     let server = server_or_skip!();
 
     // Downloader: connected to the server with its peer listener enabled.
     let listen_port = free_port().expect("free listen port");
     let mut client = Client::with_settings(server.listening_settings(
-        "e2e_downloader",
+        &format!("{peer_username}_dl"),
         "pw",
         listen_port,
     ));
@@ -741,11 +747,12 @@ fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
     let (ready_tx, ready_rx) = std::sync::mpsc::channel();
     let cfg = MockUpload {
         listen_addr: format!("127.0.0.1:{listen_port}"),
-        peer_username: "e2e_mockpeer".to_string(),
+        peer_username: peer_username.to_string(),
         filename: filename.to_string(),
         content: content.clone(),
         token,
         ready: ready_tx,
+        token_delay,
     };
     let uploader = std::thread::spawn(move || {
         if let Err(e) = run_mock_uploader(&cfg) {
@@ -764,7 +771,7 @@ fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
     let (_download, status_rx) = client
         .download(
             filename.to_string(),
-            "e2e_mockpeer".to_string(),
+            peer_username.to_string(),
             size,
             download_dir.display().to_string(),
         )
@@ -783,6 +790,16 @@ fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
     let _ = std::fs::remove_dir_all(&download_dir);
 }
 
+#[test]
+fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
+    file_downloads_over_p_and_f("e2e_mockpeer", Duration::ZERO);
+}
+
+#[test]
+fn a_file_downloads_when_the_transfer_token_arrives_late() {
+    file_downloads_over_p_and_f("e2e_latepeer", Duration::from_millis(300));
+}
+
 fn refusal_fails_the_download_quickly(
     peer_username: &str,
     build_refusal: impl FnOnce(&str) -> Message,
@@ -808,6 +825,7 @@ fn refusal_fails_the_download_quickly(
         content: Vec::new(),
         token: 0,
         ready: ready_tx,
+        token_delay: Duration::ZERO,
     };
     let uploader = std::thread::spawn(move || {
         if let Err(e) = run_refusing_uploader(&cfg, &refusal) {
@@ -905,6 +923,7 @@ fn an_interrupted_download_resumes_from_its_partial_file() {
         content: content.clone(),
         token: 424_243_u32,
         ready: ready_tx,
+        token_delay: Duration::ZERO,
     };
     let (offset_tx, offset_rx) = std::sync::mpsc::channel();
     let uploader = std::thread::spawn(move || match run_mock_uploader(&cfg) {
@@ -1122,6 +1141,7 @@ fn run_mock_direct_peer(cfg: &MockDirectUpload) -> std::io::Result<()> {
         &cfg.username,
         cfg.token,
         &cfg.content,
+        Duration::ZERO,
     )
     .map(|_| ())
 }
@@ -1269,6 +1289,7 @@ fn run_mock_firewalled_peer(cfg: &MockFirewalledUpload) -> std::io::Result<()> {
         &cfg.username,
         cfg.token,
         &cfg.content,
+        Duration::ZERO,
     )
     .map(|_| ())
 }