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
This commit is contained in:
Edison Jwa
2026-06-10 04:49:12 +09:00
parent 21fd979560
commit 6c380fc7db
2 changed files with 125 additions and 5 deletions
+121 -5
View File
@@ -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<FiletransferHandle, oneshot::Sender<Result<Vec<u8>, ProtocolError>>>;
struct EventChannels {
voice_in: mpsc::Sender<InboundVoice>,
chat: mpsc::Sender<ChatMessage>,
@@ -207,6 +212,10 @@ enum Request {
client_id: u64,
reply: oneshot::Sender<Result<ClientProfile, ProtocolError>>,
},
DownloadFile {
path: String,
reply: oneshot::Sender<Result<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<u64, Instant> = 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<Vec<u8>, 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<SocketAddr, ProtocolError> {
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<u64, Instant>,
poke_limiter: &mut PokeLimiter,
) -> Result<ClientProfile, ProtocolError> {
@@ -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<u64, Instant>,
poke_limiter: &mut PokeLimiter,
) -> Result<Vec<InMessage>, 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<u64, Instant>,
poke_limiter: &mut PokeLimiter,
) -> Result<InClientDbInfoPart, ProtocolError> {
@@ -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<Item = &'a Client>,
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
+4
View File
@@ -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),
}