Merge pull request #43 from michel/issue-41-listener-token-wait

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

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,
+    deadline: Instant,
-    peer_port: u16,
 ) {
     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.host,
-        peer_port,
+        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,
+            deadline,
-            peer_port,
         ),
         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
+/// `content`; the downloader sends an 8-byte START_DOWNLOAD offset before we
-/// the token lands in the listener's read buffer, where the download is looked
+/// send the bytes.
-/// up by token; the downloader then 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 file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
-fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
     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(|_| ())
 }