From 6c380fc7db679e56538e6a600ab2ce0ec86e5b37 Mon Sep 17 00:00:00 2001 From: Edison Jwa Date: Wed, 10 Jun 2026 04:49:12 +0900 Subject: [PATCH] feat(protocol): add file download support for avatars and icons - Add Request::DownloadFile variant with oneshot reply - Add ProtocolClient::download_avatar(client_uid) and download_icon(icon_id) - Track pending file downloads by FiletransferHandle - Handle StreamItem::FileDownload: read bytes from TCP stream - Handle StreamItem::FiletransferFailed: map to ProtocolError - Add ProtocolError::FileTransfer(String) variant - Add path helper tests for avatar/icon download paths - No tsclientlib types leak across the adapter boundary --- crates/chanora_protocol/src/adapter.rs | 126 ++++++++++++++++++++++++- crates/chanora_protocol/src/lib.rs | 4 + 2 files changed, 125 insertions(+), 5 deletions(-) diff --git a/crates/chanora_protocol/src/adapter.rs b/crates/chanora_protocol/src/adapter.rs index d5e9ee1..2341401 100644 --- a/crates/chanora_protocol/src/adapter.rs +++ b/crates/chanora_protocol/src/adapter.rs @@ -24,6 +24,7 @@ use base64::prelude::*; use chanora_resolver::ChanoraResolver; use futures::prelude::*; use std::collections::HashMap; +use tokio::io::AsyncReadExt; use tokio::sync::{mpsc, oneshot}; use tracing::{info, warn}; @@ -32,7 +33,8 @@ use tsclientlib::messages::s2c::{InClientDbInfoPart, InMessage}; use tsclientlib::prelude::*; use tsclientlib::{ ChannelId as TsChannelId, ClientId as TsClientId, Connection, ConnectionStats, - DisconnectOptions, Identity, MessageHandle, OutCommandExt, StreamItem, Version, + DisconnectOptions, FileDownloadResult, FiletransferHandle, Identity, MessageHandle, + OutCommandExt, StreamItem, Version, }; use tsproto_packets::packets::{Direction, Flags, InAudioBuf, OutCommand, OutPacket, PacketType}; use tsproto_types::ClientType; @@ -60,6 +62,9 @@ type PendingMoves = HashMap< ), >; +type PendingDownloads = + HashMap, ProtocolError>>>; + struct EventChannels { voice_in: mpsc::Sender, chat: mpsc::Sender, @@ -207,6 +212,10 @@ enum Request { client_id: u64, reply: oneshot::Sender>, }, + DownloadFile { + path: String, + reply: oneshot::Sender, ProtocolError>>, + }, } /// Why a [`ProtocolClient`] task ended. Distinguishes a user-driven @@ -366,6 +375,26 @@ impl ProtocolClient { .map_err(|_| ProtocolError::Lost("client_profile reply dropped".to_string()))? } + async fn download_file(&self, path: String) -> Result, ProtocolError> { + let (tx, rx) = oneshot::channel(); + self.tx + .send(Request::DownloadFile { path, reply: tx }) + .await + .map_err(|_| ProtocolError::Lost("connection task is gone".to_string()))?; + rx.await + .map_err(|_| ProtocolError::Lost("download_file reply dropped".to_string()))? + } + + /// Download the current avatar bytes for a TeamSpeak client UID. + pub async fn download_avatar(&self, client_uid: &str) -> Result, ProtocolError> { + self.download_file(avatar_download_path(client_uid)).await + } + + /// Download the current channel/server icon bytes for an icon id. + pub async fn download_icon(&self, icon_id: u64) -> Result, ProtocolError> { + self.download_file(icon_download_path(icon_id)).await + } + /// Disconnect cleanly. Blocks until the task exits. pub async fn disconnect(self) { let (tx, rx) = oneshot::channel(); @@ -713,6 +742,7 @@ async fn connection_task( // deadline so a server that never replies doesn't leak the // reply channel — at most 3 s of pending state per move. let mut pending_moves: PendingMoves = HashMap::new(); + let mut pending_downloads: PendingDownloads = HashMap::new(); let mut voice_activity: HashMap = HashMap::new(); let mut poke_limiter = PokeLimiter::new(); @@ -737,6 +767,12 @@ async fn connection_task( StreamItem::Audio(buf) => { handle_audio_stream_item(&channels.voice_in, &mut voice_activity, buf).await; } + StreamItem::FileDownload(handle, result) => { + handle_download_stream_item(&mut pending_downloads, handle, result).await; + } + StreamItem::FiletransferFailed(handle, error) => { + handle_download_failure(&mut pending_downloads, handle, error); + } other => handle_non_audio_stream_item( &con, other, @@ -843,12 +879,25 @@ async fn connection_task( client_id, &channels, &mut pending_moves, + &mut pending_downloads, &mut voice_activity, &mut poke_limiter, ) .await; let _ = reply.send(r); } + Ok(Request::DownloadFile { path, reply }) => { + match con.download_file(TsChannelId(0), &path, None, None) { + Ok(handle) => { + pending_downloads.insert(handle, reply); + } + Err(e) => { + let _ = reply.send(Err(ProtocolError::FileTransfer(format!( + "start download {path}: {e}" + )))); + } + } + } Ok(Request::Disconnect(reply)) => { let _ = con.disconnect(DisconnectOptions::new()); bounded_drain_stream(con.events(), DISCONNECT_EVENT_DRAIN_TIMEOUT).await; @@ -1024,6 +1073,39 @@ fn handle_non_audio_stream_item( } } +async fn handle_download_stream_item( + pending_downloads: &mut PendingDownloads, + handle: FiletransferHandle, + result: FileDownloadResult, +) { + if let Some(reply) = pending_downloads.remove(&handle) { + let _ = reply.send(read_download_bytes(result).await); + } +} + +fn handle_download_failure( + pending_downloads: &mut PendingDownloads, + handle: FiletransferHandle, + error: tsclientlib::Error, +) { + if let Some(reply) = pending_downloads.remove(&handle) { + let _ = reply.send(Err(ProtocolError::FileTransfer(error.to_string()))); + } +} + +async fn read_download_bytes(result: FileDownloadResult) -> Result, ProtocolError> { + let size = usize::try_from(result.size).map_err(|_| { + ProtocolError::FileTransfer(format!("download too large to buffer: {} bytes", result.size)) + })?; + let mut stream = result.stream; + let mut bytes = vec![0_u8; size]; + stream + .read_exact(&mut bytes) + .await + .map_err(|e| ProtocolError::FileTransfer(e.to_string()))?; + Ok(bytes) +} + async fn resolve_server_socket(address: &str) -> Result { let resolver = ChanoraResolver::new().map_err(|err| ProtocolError::DnsFailed { host: address.to_string(), @@ -1174,6 +1256,7 @@ async fn fetch_client_profile( client_id: u64, channels: &EventChannels, pending_moves: &mut PendingMoves, + pending_downloads: &mut PendingDownloads, voice_activity: &mut HashMap, poke_limiter: &mut PokeLimiter, ) -> Result { @@ -1220,6 +1303,7 @@ async fn fetch_client_profile( build_command("servergrouplist", &[], &[]), channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1231,6 +1315,7 @@ async fn fetch_client_profile( build_command("channelgrouplist", &[], &[]), channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1246,6 +1331,7 @@ async fn fetch_client_profile( ), channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1265,6 +1351,7 @@ async fn fetch_client_profile( build_command("getconnectioninfo", &[("clid", client_id.to_string())], &[]), channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1285,6 +1372,7 @@ async fn fetch_client_profile( database_id, channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1448,6 +1536,7 @@ async fn request_messages( command: OutCommand, channels: &EventChannels, pending_moves: &mut PendingMoves, + pending_downloads: &mut PendingDownloads, voice_activity: &mut HashMap, poke_limiter: &mut PokeLimiter, ) -> Result, ProtocolError> { @@ -1485,6 +1574,12 @@ async fn request_messages( StreamItem::Audio(buf) => { handle_audio_stream_item(&channels.voice_in, voice_activity, buf).await; } + StreamItem::FileDownload(handle, result) => { + handle_download_stream_item(pending_downloads, handle, result).await; + } + StreamItem::FiletransferFailed(handle, error) => { + handle_download_failure(pending_downloads, handle, error); + } other => handle_non_audio_stream_item( con, other, @@ -1503,6 +1598,7 @@ async fn request_client_db_info( dbid: tsclientlib::ClientDbId, channels: &EventChannels, pending_moves: &mut PendingMoves, + pending_downloads: &mut PendingDownloads, voice_activity: &mut HashMap, poke_limiter: &mut PokeLimiter, ) -> Result { @@ -1511,6 +1607,7 @@ async fn request_client_db_info( build_command("clientdbinfo", &[("cldbid", dbid.0.to_string())], &[]), channels, pending_moves, + pending_downloads, voice_activity, poke_limiter, ) @@ -1568,6 +1665,14 @@ fn uid_to_avatar_path(uid_b64: &str) -> String { rendered } +fn avatar_download_path(client_uid: &str) -> String { + format!("/avatar_{}", uid_to_avatar_path(client_uid)) +} + +fn icon_download_path(icon_id: u64) -> String { + format!("/icon_{icon_id}") +} + fn find_client_by_id<'a>( clients: impl IntoIterator, client_id: u64, @@ -1951,10 +2056,11 @@ const _: () = { #[cfg(test)] mod tests { use super::{ - bounded_drain_stream, client_profile_refresh_plan, drain_voice_packets_for_tick, - is_server_query_client_type, send_with_timeout, server_socket_from_config, - sort_channels_tree_by, std_duration_millis, ConnectConfig, ProtocolClient, Request, - SendTimeoutError, DISCONNECT_REPLY_TIMEOUT, + avatar_download_path, bounded_drain_stream, client_profile_refresh_plan, + drain_voice_packets_for_tick, icon_download_path, is_server_query_client_type, + send_with_timeout, server_socket_from_config, sort_channels_tree_by, + std_duration_millis, ConnectConfig, ProtocolClient, Request, SendTimeoutError, + DISCONNECT_REPLY_TIMEOUT, }; use futures::stream; use std::time::Duration; @@ -2052,6 +2158,16 @@ mod tests { assert!(plan.needs_channel_groups); } + #[test] + fn avatar_download_path_uses_uid_hex_encoding() { + assert_eq!(avatar_download_path("AQID"), "/avatar_abacad"); + } + + #[test] + fn icon_download_path_uses_unsigned_icon_id() { + assert_eq!(icon_download_path(42), "/icon_42"); + } + #[test] fn channel_sort_linked_list_under_one_parent() { // Server emits four root-level channels in arbitrary HashMap diff --git a/crates/chanora_protocol/src/lib.rs b/crates/chanora_protocol/src/lib.rs index 3d5d409..4d77ce6 100644 --- a/crates/chanora_protocol/src/lib.rs +++ b/crates/chanora_protocol/src/lib.rs @@ -115,4 +115,8 @@ pub enum ProtocolError { /// should never see this; if they do, it is a mapping bug here. #[error("protocol backend: {0}")] Backend(String), + + /// A file transfer failed while downloading protocol-owned assets. + #[error("file transfer failed: {0}")] + FileTransfer(String), }