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:
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user