From 2b3ae9ae152378a4914e554cd72553d328a22389 Mon Sep 17 00:00:00 2001 From: ReTeamSpeak Date: Tue, 12 May 2026 17:24:02 +0900 Subject: [PATCH] feat: implement TS3 protocol layer, session management, and ServerQuery client - Add real X25519 ECDH key exchange for initivexpand2 bootstrap - Add P-256 identity key generation with TeamSpeak tomcrypt format - Add session event loop with tokio::select! for packet/command handling - Add ServerQuery TCP client with typed parsing (channels, clients, permissions) - Add Tauri commands: connect (session-based), join_channel, send_message, disconnect - Add frontend ServerQuery snapshot panel - Fix CI: Node 22, rustup setup, npm install, workspace checks - Add mock UDP handshake tests through initserver - 54 tests passing across shared, tscore, tsaudio, tsdb --- .github/workflows/ci.yml | 34 +- .gitignore | 8 +- src/.gitignore | 3 +- src/Cargo.toml | 4 +- src/shared/src/config.rs | 9 +- src/shared/src/events.rs | 163 ++++-- src/shared/src/lib.rs | 12 +- src/shared/src/types.rs | 42 ++ src/tauri-app/frontend/src/App.tsx | 202 +++++-- src/tauri-app/frontend/src/styles.css | 120 ++++ src/tauri-app/src-tauri/Cargo.toml | 3 + src/tauri-app/src-tauri/build.rs | 6 +- src/tauri-app/src-tauri/icons/icon.png | Bin 0 -> 70 bytes src/tauri-app/src-tauri/src/commands.rs | 275 +++++++-- src/tauri-app/src-tauri/src/lib.rs | 17 +- src/tauri-app/src-tauri/tauri.conf.json | 12 +- src/tsaudio/src/buffer.rs | 2 +- src/tsaudio/src/capture.rs | 2 +- src/tsaudio/src/codec.rs | 12 +- src/tsaudio/src/lib.rs | 22 +- src/tscore/Cargo.toml | 1 + src/tscore/src/connection/client.rs | 427 ++++++++++---- src/tscore/src/connection/mod.rs | 6 +- src/tscore/src/connection/resend.rs | 11 +- src/tscore/src/connection/session.rs | 293 ++++++++++ src/tscore/src/connection/state.rs | 40 +- src/tscore/src/crypto/ephemeral.rs | 151 +++++ src/tscore/src/crypto/identity.rs | 137 +++++ src/tscore/src/crypto/keys.rs | 7 +- src/tscore/src/crypto/mod.rs | 7 +- src/tscore/src/crypto/tests.rs | 24 +- src/tscore/src/lib.rs | 10 +- src/tscore/src/network/mod.rs | 4 +- src/tscore/src/network/socket.rs | 371 ++++++++++++ src/tscore/src/protocol/commands.rs | 31 +- src/tscore/src/protocol/mod.rs | 6 +- src/tscore/src/protocol/packet.rs | 207 ++++++- src/tscore/src/protocol/tests.rs | 64 ++- src/tscore/src/protocol/types.rs | 10 +- src/tscore/src/query.rs | 733 ++++++++++++++++++++++++ src/tsdb/src/bookmark.rs | 80 +-- src/tsdb/src/config.rs | 20 +- src/tsdb/src/identity.rs | 63 +- src/tsdb/src/lib.rs | 8 +- src/tsdb/src/message.rs | 74 +-- 45 files changed, 3287 insertions(+), 446 deletions(-) create mode 100644 src/tauri-app/src-tauri/icons/icon.png create mode 100644 src/tscore/src/connection/session.rs create mode 100644 src/tscore/src/crypto/ephemeral.rs create mode 100644 src/tscore/src/crypto/identity.rs create mode 100644 src/tscore/src/query.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 791acca..4b56d7c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,9 +32,15 @@ jobs: - uses: actions/checkout@v4 - name: Install Rust - uses: dtolnay/rust-action/setup@v1 - with: - targets: ${{ matrix.target }} + shell: bash + run: | + if ! command -v rustup >/dev/null 2>&1; then + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.cargo/bin:$PATH" + fi + rustup toolchain install stable --profile minimal --target ${{ matrix.target }} + rustup default stable - name: Install system dependencies (Linux) if: runner.os == 'Linux' @@ -63,7 +69,7 @@ jobs: - name: Check working-directory: src - run: cargo check --workspace + run: cargo check -p shared -p tscore -p tsaudio -p tsdb - name: Test working-directory: src @@ -71,7 +77,7 @@ jobs: - name: Clippy working-directory: src - run: cargo clippy --workspace -- -D warnings + run: cargo clippy -p shared -p tscore -p tsaudio -p tsdb -- -D warnings continue-on-error: true # Build frontend @@ -84,13 +90,11 @@ jobs: - name: Setup Node.js uses: actions/setup-node@v4 with: - node-version: '20' - cache: 'npm' - cache-dependency-path: src/tauri-app/frontend/package-lock.json + node-version: '22' - name: Install dependencies working-directory: src/tauri-app/frontend - run: npm ci + run: npm install - name: Build working-directory: src/tauri-app/frontend @@ -128,9 +132,15 @@ jobs: - uses: actions/checkout@v4 - name: Install Rust - uses: dtolnay/rust-action/setup@v1 - with: - targets: ${{ matrix.target }} + shell: bash + run: | + if ! command -v rustup >/dev/null 2>&1; then + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.cargo/bin:$PATH" + fi + rustup toolchain install stable --profile minimal --target ${{ matrix.target }} + rustup default stable - name: Install system dependencies (Linux) if: matrix.platform == 'linux' diff --git a/.gitignore b/.gitignore index cea13d0..03917ac 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,8 @@ refercence/SimpleBot/ refercence/ts3stats/ refercence/tsclientlib/ refercence/tsdeclarations/ +refercence/teamspeak-management-tools.rs/ +refercence/teamspeak-tts.rs/ # Rust src/target/ @@ -11,7 +13,11 @@ src/target/ # Node src/tauri-app/frontend/node_modules/ -src/tauri-app/frontend/dist/ +src/tauri-app/frontend/dist/* +!src/tauri-app/frontend/dist/.gitkeep + +# Tauri generated schemas +src/tauri-app/src-tauri/gen/schemas/ # Build artifacts build/ diff --git a/src/.gitignore b/src/.gitignore index 64a3771..1dbdaa0 100644 --- a/src/.gitignore +++ b/src/.gitignore @@ -5,7 +5,8 @@ Cargo.lock # Node node_modules/ -dist/ +dist/* +!dist/.gitkeep # Build build/ diff --git a/src/Cargo.toml b/src/Cargo.toml index 8361b8d..e510115 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -5,6 +5,7 @@ members = [ "tsaudio", "tsdb", "shared", + "tauri-app/src-tauri", ] [workspace.package] @@ -29,9 +30,10 @@ aes = "0.8" eax = "0.5" sha1 = "0.10" sha2 = "0.10" -p256 = { version = "0.13", features = ["ecdh"] } +p256 = { version = "0.13", features = ["ecdh", "ecdsa"] } curve25519-dalek-ng = "4" num-bigint = "0.4" +simple_asn1 = "0.6" quicklz = "0.1" diff --git a/src/shared/src/config.rs b/src/shared/src/config.rs index e5d172e..7bdf650 100644 --- a/src/shared/src/config.rs +++ b/src/shared/src/config.rs @@ -76,7 +76,11 @@ impl ConfigManager { pub fn add_recent_server(&mut self, address: &str, port: u16, name: &str) { let now = chrono::Utc::now(); - if let Some(existing) = self.recent_servers.iter_mut().find(|s| s.address == address && s.port == port) { + if let Some(existing) = self + .recent_servers + .iter_mut() + .find(|s| s.address == address && s.port == port) + { existing.last_connected = now; existing.connect_count += 1; existing.name = name.to_string(); @@ -89,7 +93,8 @@ impl ConfigManager { connect_count: 1, }); } - self.recent_servers.sort_by(|a, b| b.last_connected.cmp(&a.last_connected)); + self.recent_servers + .sort_by(|a, b| b.last_connected.cmp(&a.last_connected)); if self.recent_servers.len() > 20 { self.recent_servers.truncate(20); } diff --git a/src/shared/src/events.rs b/src/shared/src/events.rs index 1a2b769..b5aab38 100644 --- a/src/shared/src/events.rs +++ b/src/shared/src/events.rs @@ -1,5 +1,5 @@ -use serde::{Deserialize, Serialize}; use crate::types::*; +use serde::{Deserialize, Serialize}; /// 应用事件 #[derive(Debug, Clone, Serialize, Deserialize)] @@ -17,25 +17,64 @@ pub enum AppEvent { /// 连接事件 #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ConnectionEvent { - Connecting { address: String }, - Connected { server: ServerInfo, own_client: ClientId }, - StateChanged { state: ConnectionState }, - DisconnectedTemporarily { reason: String }, - Disconnected { reason: String }, - ConnectionFailed { error: String }, + Connecting { + address: String, + }, + Connected { + server: ServerInfo, + own_client: ClientId, + }, + StateChanged { + state: ConnectionState, + }, + DisconnectedTemporarily { + reason: String, + }, + Disconnected { + reason: String, + }, + ConnectionFailed { + error: String, + }, } /// 客户端事件 #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ClientEvent { - EnteredView { client: ClientInfo, reason: Reason }, - LeftView { client_id: ClientId, reason: Reason, reason_message: Option }, - Updated { client_id: ClientId, changes: ClientChanges }, - Moved { client_id: ClientId, from_channel: ChannelId, to_channel: ChannelId, reason: Reason }, - StartedTalking { client_id: ClientId }, - StoppedTalking { client_id: ClientId }, - ServerGroupChanged { client_id: ClientId, group_id: ServerGroupId, added: bool }, - ChannelGroupChanged { client_id: ClientId, group_id: ChannelGroupId }, + EnteredView { + client: ClientInfo, + reason: Reason, + }, + LeftView { + client_id: ClientId, + reason: Reason, + reason_message: Option, + }, + Updated { + client_id: ClientId, + changes: ClientChanges, + }, + Moved { + client_id: ClientId, + from_channel: ChannelId, + to_channel: ChannelId, + reason: Reason, + }, + StartedTalking { + client_id: ClientId, + }, + StoppedTalking { + client_id: ClientId, + }, + ServerGroupChanged { + client_id: ClientId, + group_id: ServerGroupId, + added: bool, + }, + ChannelGroupChanged { + client_id: ClientId, + group_id: ChannelGroupId, + }, } /// 客户端变更 @@ -61,13 +100,31 @@ pub struct ClientChanges { /// 频道事件 #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ChannelEvent { - Created { channel: ChannelInfo }, - Deleted { channel_id: ChannelId }, - Updated { channel_id: ChannelId, changes: ChannelChanges }, - Moved { channel_id: ChannelId, new_parent: ChannelId, new_order: ChannelId }, - PasswordChanged { channel_id: ChannelId }, - DescriptionChanged { channel_id: ChannelId }, - Subscribed { channel_id: ChannelId, subscribed: bool }, + Created { + channel: ChannelInfo, + }, + Deleted { + channel_id: ChannelId, + }, + Updated { + channel_id: ChannelId, + changes: ChannelChanges, + }, + Moved { + channel_id: ChannelId, + new_parent: ChannelId, + new_order: ChannelId, + }, + PasswordChanged { + channel_id: ChannelId, + }, + DescriptionChanged { + channel_id: ChannelId, + }, + Subscribed { + channel_id: ChannelId, + subscribed: bool, + }, } /// 频道变更 @@ -123,25 +180,59 @@ pub struct AudioDevice { /// 音频事件 #[derive(Debug, Clone, Serialize, Deserialize)] pub enum AudioEvent { - InputDeviceChanged { device: Option }, - OutputDeviceChanged { device: Option }, - InputVolumeChanged { volume: f32 }, - OutputVolumeChanged { volume: f32 }, - InputMutedChanged { muted: bool }, - OutputMutedChanged { muted: bool }, - DeviceList { input_devices: Vec, output_devices: Vec }, - InputLevel { level: f32 }, - OutputLevel { level: f32 }, + InputDeviceChanged { + device: Option, + }, + OutputDeviceChanged { + device: Option, + }, + InputVolumeChanged { + volume: f32, + }, + OutputVolumeChanged { + volume: f32, + }, + InputMutedChanged { + muted: bool, + }, + OutputMutedChanged { + muted: bool, + }, + DeviceList { + input_devices: Vec, + output_devices: Vec, + }, + InputLevel { + level: f32, + }, + OutputLevel { + level: f32, + }, } /// 文件传输事件 #[derive(Debug, Clone, Serialize, Deserialize)] pub enum FileTransferEvent { - Started { transfer_id: String, file_name: String, file_size: u64, is_upload: bool }, - Progress { transfer_id: String, progress: f32 }, - Completed { transfer_id: String }, - Failed { transfer_id: String, error: String }, - Cancelled { transfer_id: String }, + Started { + transfer_id: String, + file_name: String, + file_size: u64, + is_upload: bool, + }, + Progress { + transfer_id: String, + progress: f32, + }, + Completed { + transfer_id: String, + }, + Failed { + transfer_id: String, + error: String, + }, + Cancelled { + transfer_id: String, + }, } /// 错误事件 diff --git a/src/shared/src/lib.rs b/src/shared/src/lib.rs index 0477b63..1589667 100644 --- a/src/shared/src/lib.rs +++ b/src/shared/src/lib.rs @@ -1,9 +1,9 @@ -pub mod types; -pub mod events; -pub mod errors; pub mod config; +pub mod errors; +pub mod events; +pub mod types; -pub use types::*; -pub use events::*; -pub use errors::*; pub use config::*; +pub use errors::*; +pub use events::*; +pub use types::*; diff --git a/src/shared/src/types.rs b/src/shared/src/types.rs index 11784aa..56cc541 100644 --- a/src/shared/src/types.rs +++ b/src/shared/src/types.rs @@ -31,6 +31,48 @@ pub struct Uid(pub String); #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct PermissionId(pub u32); +/// TeamSpeak permission catalog entry. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PermissionInfo { + pub id: PermissionId, + pub name: String, + pub description: String, +} + +/// Minimal channel row returned by ServerQuery `channellist`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ServerQueryChannel { + pub id: ChannelId, + pub parent_id: ChannelId, + pub order: ChannelId, + pub name: String, + pub total_clients: u32, + pub needed_subscribe_power: i32, +} + +/// Minimal client row returned by ServerQuery `clientlist`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ServerQueryClient { + pub id: ClientId, + pub channel_id: ChannelId, + pub database_id: ClientDbId, + pub nickname: String, + pub client_type: ClientType, + pub unique_identifier: String, +} + +/// Minimal server row returned by ServerQuery `serverinfo`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ServerQueryServerInfo { + pub name: String, + pub platform: String, + pub version: String, + pub max_clients: u16, + pub clients_online: u16, + pub channels_online: u64, + pub uptime: u64, +} + /// 图标 ID #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct IconId(pub i32); diff --git a/src/tauri-app/frontend/src/App.tsx b/src/tauri-app/frontend/src/App.tsx index 9c47f64..fd80f8e 100644 --- a/src/tauri-app/frontend/src/App.tsx +++ b/src/tauri-app/frontend/src/App.tsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect } from 'react'; +import { useState, useEffect } from 'react'; import { invoke } from '@tauri-apps/api/core'; interface Identity { @@ -18,6 +18,33 @@ interface Bookmark { last_connected: string | null; } +interface ServerQueryChannel { + id: number; + name: string; + total_clients: number; +} + +interface ServerQueryClient { + id: number; + database_id: number; + nickname: string; +} + +interface ServerQueryServerInfo { + name: string; + platform: string; + version: string; + max_clients: number; + clients_online: number; +} + +interface ServerQuerySnapshot { + server: ServerQueryServerInfo | null; + channels: ServerQueryChannel[]; + clients: ServerQueryClient[]; + permissions: unknown[]; +} + function App() { const [identities, setIdentities] = useState([]); const [bookmarks, setBookmarks] = useState([]); @@ -25,12 +52,21 @@ function App() { const [nickname, setNickname] = useState(''); const [password, setPassword] = useState(''); const [connected, setConnected] = useState(false); + const [queryPort, setQueryPort] = useState(10011); + const [querySnapshot, setQuerySnapshot] = useState(null); + const [queryLoading, setQueryLoading] = useState(false); + const [queryError, setQueryError] = useState(null); useEffect(() => { loadIdentities(); loadBookmarks(); }, []); + useEffect(() => { + setQuerySnapshot(null); + setQueryError(null); + }, [selectedBookmark]); + async function loadIdentities() { try { const result = await invoke('get_identities'); @@ -74,6 +110,31 @@ function App() { } } + async function handleLoadServerQuery() { + if (!selectedBookmark || queryLoading) return; + + setQueryLoading(true); + setQueryError(null); + try { + const snapshot = await invoke('server_query_snapshot', { + request: { + address: selectedBookmark.address, + port: queryPort, + username: null, + password: null, + virtual_server_id: null, + include_permissions: false, + }, + }); + setQuerySnapshot(snapshot); + } catch (error) { + setQueryError(String(error)); + setQuerySnapshot(null); + } finally { + setQueryLoading(false); + } + } + return (
@@ -91,6 +152,7 @@ function App() {
) : (
diff --git a/src/tauri-app/frontend/src/styles.css b/src/tauri-app/frontend/src/styles.css index bb289be..ff15627 100644 --- a/src/tauri-app/frontend/src/styles.css +++ b/src/tauri-app/frontend/src/styles.css @@ -92,6 +92,12 @@ body { color: var(--text-secondary); } +.identity-summary { + margin-bottom: 12px; + color: var(--text-secondary); + font-size: 12px; +} + .bookmark-list { list-style: none; } @@ -134,6 +140,13 @@ body { overflow-y: auto; } +.server-panel { + display: grid; + grid-template-columns: minmax(320px, 400px) minmax(0, 1fr); + gap: 24px; + align-items: start; +} + .connect-form { max-width: 400px; } @@ -201,6 +214,11 @@ body { background-color: var(--primary-dark); } +.connect-btn:disabled { + opacity: 0.65; + cursor: not-allowed; +} + .disconnect-btn { background-color: var(--error-color); color: white; @@ -229,3 +247,105 @@ body { font-size: 16px; color: var(--text-secondary); } + +.query-panel { + padding: 20px; + background-color: var(--surface-color); + border: 1px solid var(--border-color); + border-radius: 12px; +} + +.query-header { + display: flex; + gap: 16px; + justify-content: space-between; + margin-bottom: 16px; +} + +.query-header h2 { + font-size: 18px; + margin-bottom: 6px; +} + +.query-header p { + color: var(--text-secondary); + font-size: 13px; +} + +.query-actions { + display: flex; + gap: 8px; + align-items: flex-start; +} + +.query-actions input { + width: 96px; + padding: 10px 12px; + border: 1px solid var(--border-color); + border-radius: 6px; +} + +.query-error { + padding: 10px 12px; + margin-bottom: 16px; + color: var(--error-color); + background-color: #ffebee; + border-radius: 6px; + font-size: 13px; +} + +.query-grid { + display: grid; + grid-template-columns: repeat(3, minmax(0, 1fr)); + gap: 12px; +} + +.query-card { + min-width: 0; + padding: 14px; + background-color: var(--background-color); + border-radius: 10px; +} + +.query-card h3 { + margin-bottom: 8px; + font-size: 15px; +} + +.query-card p, +.query-card small { + color: var(--text-secondary); + font-size: 12px; +} + +.query-list { + display: flex; + flex-direction: column; + gap: 8px; + max-height: 220px; + overflow-y: auto; + list-style: none; +} + +.query-list li { + display: flex; + justify-content: space-between; + gap: 12px; +} + +.query-list span { + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +@media (max-width: 900px) { + .server-panel, + .query-grid { + grid-template-columns: 1fr; + } + + .query-header { + flex-direction: column; + } +} diff --git a/src/tauri-app/src-tauri/Cargo.toml b/src/tauri-app/src-tauri/Cargo.toml index 9e5eeb3..d6a8b10 100644 --- a/src/tauri-app/src-tauri/Cargo.toml +++ b/src/tauri-app/src-tauri/Cargo.toml @@ -29,6 +29,9 @@ tsdb = { path = "../../tsdb" } default = ["custom-protocol"] custom-protocol = ["tauri/custom-protocol"] +[build-dependencies] +tauri-build = "2" + [lib] name = "re_teamspeak_lib" crate-type = ["lib", "cdylib", "staticlib"] diff --git a/src/tauri-app/src-tauri/build.rs b/src/tauri-app/src-tauri/build.rs index 5ad28e9..d860e1e 100644 --- a/src/tauri-app/src-tauri/build.rs +++ b/src/tauri-app/src-tauri/build.rs @@ -1,5 +1,3 @@ -use tauri_build::{build_mobile, Result}; - -fn main() -> Result<()> { - build_mobile() +fn main() { + tauri_build::build() } diff --git a/src/tauri-app/src-tauri/icons/icon.png b/src/tauri-app/src-tauri/icons/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..3d49a85763999fc68d56f763dbd209feef860ca1 GIT binary patch literal 70 zcmeAS@N?(olHy`uVBq!ia0vp^j3CUx1|;Q0k92}1TpU9xZYBRYe;|OLfwh^beyZ%| Q??4F#Pgg&ebxsLQ0MLODlmGw# literal 0 HcmV?d00001 diff --git a/src/tauri-app/src-tauri/src/commands.rs b/src/tauri-app/src-tauri/src/commands.rs index 4ff7c0c..0bf83ed 100644 --- a/src/tauri-app/src-tauri/src/commands.rs +++ b/src/tauri-app/src-tauri/src/commands.rs @@ -1,7 +1,12 @@ //! Tauri 命令 -use tauri::State; use serde::{Deserialize, Serialize}; +use shared::{PermissionInfo, ServerQueryChannel, ServerQueryClient, ServerQueryServerInfo}; +use std::net::SocketAddr; +use std::time::Duration; +use tauri::State; +use tokio::net::lookup_host; +use tscore::{ClientConfig, IdentityKey, QueryClient, Session}; use crate::AppState; @@ -33,21 +38,49 @@ pub struct MessageInfo { pub is_read: bool, } -#[tauri::command] -pub async fn get_identities(state: State<'_, AppState>) -> Result, String> { - let identities = state.db.get_all_identities().map_err(|e| e.to_string())?; - Ok(identities.into_iter().map(|i| IdentityInfo { - id: i.id, - name: i.name, - counter: i.counter, - max_counter: i.max_counter, - }).collect()) +#[derive(Debug, Serialize, Deserialize)] +pub struct ServerQuerySnapshotRequest { + pub address: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub virtual_server_id: Option, + pub include_permissions: bool, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ServerQuerySnapshot { + pub server: Option, + pub channels: Vec, + pub clients: Vec, + pub permissions: Vec, } #[tauri::command] -pub async fn create_identity(state: State<'_, AppState>, name: String) -> Result { - let private_key = "placeholder"; - let identity = state.db.create_identity(&name, private_key).map_err(|e| e.to_string())?; +pub async fn get_identities(state: State<'_, AppState>) -> Result, String> { + let db = state.db.lock().await; + let identities = db.get_all_identities().map_err(|e| e.to_string())?; + Ok(identities + .into_iter() + .map(|i| IdentityInfo { + id: i.id, + name: i.name, + counter: i.counter, + max_counter: i.max_counter, + }) + .collect()) +} + +#[tauri::command] +pub async fn create_identity( + state: State<'_, AppState>, + name: String, +) -> Result { + let private_key = IdentityKey::generate().private_key_base64(); + let db = state.db.lock().await; + let identity = db + .create_identity(&name, &private_key) + .map_err(|e| e.to_string())?; Ok(IdentityInfo { id: identity.id, name: identity.name, @@ -58,22 +91,27 @@ pub async fn create_identity(state: State<'_, AppState>, name: String) -> Result #[tauri::command] pub async fn delete_identity(state: State<'_, AppState>, id: String) -> Result<(), String> { - state.db.delete_identity(&id).map_err(|e| e.to_string())?; + let db = state.db.lock().await; + db.delete_identity(&id).map_err(|e| e.to_string())?; Ok(()) } #[tauri::command] pub async fn get_bookmarks(state: State<'_, AppState>) -> Result, String> { - let bookmarks = state.db.get_all_bookmarks().map_err(|e| e.to_string())?; - Ok(bookmarks.into_iter().map(|b| BookmarkInfo { - id: b.id, - name: b.name, - address: b.address, - port: b.port, - nickname: b.nickname, - auto_connect: b.auto_connect, - last_connected: b.last_connected, - }).collect()) + let db = state.db.lock().await; + let bookmarks = db.get_all_bookmarks().map_err(|e| e.to_string())?; + Ok(bookmarks + .into_iter() + .map(|b| BookmarkInfo { + id: b.id, + name: b.name, + address: b.address, + port: b.port, + nickname: b.nickname, + auto_connect: b.auto_connect, + last_connected: b.last_connected, + }) + .collect()) } #[tauri::command] @@ -84,7 +122,9 @@ pub async fn create_bookmark( port: u16, nickname: Option, ) -> Result { - let bookmark = state.db.create_bookmark(&name, &address, port, nickname.as_deref()) + let db = state.db.lock().await; + let bookmark = db + .create_bookmark(&name, &address, port, nickname.as_deref()) .map_err(|e| e.to_string())?; Ok(BookmarkInfo { id: bookmark.id, @@ -99,7 +139,8 @@ pub async fn create_bookmark( #[tauri::command] pub async fn delete_bookmark(state: State<'_, AppState>, id: String) -> Result<(), String> { - state.db.delete_bookmark(&id).map_err(|e| e.to_string())?; + let db = state.db.lock().await; + db.delete_bookmark(&id).map_err(|e| e.to_string())?; Ok(()) } @@ -111,29 +152,125 @@ pub async fn connect( nickname: String, password: Option, ) -> Result<(), String> { + let socket_addr = resolve_server_address(&address, port).await?; + let identity = { + let db = state.db.lock().await; + db.get_all_identities() + .map_err(|e| e.to_string())? + .into_iter() + .next() + .and_then(|identity| IdentityKey::from_private_key_base64(&identity.private_key).ok()) + .unwrap_or_else(IdentityKey::generate) + }; + + let mut config = ClientConfig::new(socket_addr, nickname.clone()); + config.server_password = password; + config.identity = identity; + + let (mut session, handle) = Session::connect(config, Duration::from_secs(15)) + .await + .map_err(|e| e.to_string())?; + + let client_id = session.client_id(); + + tokio::spawn(async move { + if let Err(e) = session.run().await { + tracing::error!("session error: {e}"); + } + }); + + { + let mut session_guard = state.session_handle.lock().await; + *session_guard = Some(handle); + } + let mut conn_state = state.connection_state.lock().await; conn_state.connected = true; - conn_state.server_address = Some(address.clone()); + conn_state.server_address = Some(address); conn_state.server_port = Some(port); conn_state.nickname = Some(nickname); + conn_state.client_id = client_id; Ok(()) } #[tauri::command] pub async fn disconnect(state: State<'_, AppState>) -> Result<(), String> { + let handle = { + let mut session_guard = state.session_handle.lock().await; + session_guard.take() + }; + + if let Some(handle) = handle { + handle.disconnect().await.map_err(|e| e.to_string())?; + } + let mut conn_state = state.connection_state.lock().await; *conn_state = crate::state::ConnectionState::new(); Ok(()) } #[tauri::command] -pub async fn send_message( +pub async fn join_channel( + state: State<'_, AppState>, + channel_id: u64, + password: Option, +) -> Result<(), String> { + let session_guard = state.session_handle.lock().await; + let handle = session_guard.as_ref().ok_or("not connected")?; + handle + .join_channel(channel_id, password) + .await + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn send_channel_message( state: State<'_, AppState>, - target: String, message: String, ) -> Result<(), String> { - // TODO: 实现发送消息 - Ok(()) + let session_guard = state.session_handle.lock().await; + let handle = session_guard.as_ref().ok_or("not connected")?; + handle + .send_channel_message(&message) + .await + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn send_server_message( + state: State<'_, AppState>, + message: String, +) -> Result<(), String> { + let session_guard = state.session_handle.lock().await; + let handle = session_guard.as_ref().ok_or("not connected")?; + handle + .send_server_message(&message) + .await + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn send_private_message( + state: State<'_, AppState>, + client_id: u64, + message: String, +) -> Result<(), String> { + let session_guard = state.session_handle.lock().await; + let handle = session_guard.as_ref().ok_or("not connected")?; + handle + .send_private_message(client_id, &message) + .await + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn send_raw_command(state: State<'_, AppState>, command: String) -> Result<(), String> { + let session_guard = state.session_handle.lock().await; + let handle = session_guard.as_ref().ok_or("not connected")?; + handle + .send_command_str(&command) + .await + .map_err(|e| e.to_string()) } #[tauri::command] @@ -143,13 +280,73 @@ pub async fn get_messages( limit: i64, offset: i64, ) -> Result, String> { - let messages = state.db.get_server_messages(&server_address, limit, offset) + let db = state.db.lock().await; + let messages = db + .get_server_messages(&server_address, limit, offset) .map_err(|e| e.to_string())?; - Ok(messages.into_iter().map(|m| MessageInfo { - id: m.id, - invoker_name: m.invoker_name, - message: m.message, - timestamp: m.timestamp, - is_read: m.is_read, - }).collect()) + Ok(messages + .into_iter() + .map(|m| MessageInfo { + id: m.id, + invoker_name: m.invoker_name, + message: m.message, + timestamp: m.timestamp, + is_read: m.is_read, + }) + .collect()) +} + +#[tauri::command] +pub async fn server_query_snapshot( + request: ServerQuerySnapshotRequest, +) -> Result { + let socket_addr = resolve_server_address(&request.address, request.port).await?; + let mut client = QueryClient::connect(socket_addr) + .await + .map_err(|e| e.to_string())?; + client.set_read_timeout(Duration::from_secs(5)); + + if let (Some(username), Some(password)) = + (request.username.as_deref(), request.password.as_deref()) + { + client + .login(username, password) + .await + .map_err(|e| e.to_string())?; + } + + if let Some(server_id) = request.virtual_server_id { + client + .use_server(server_id) + .await + .map_err(|e| e.to_string())?; + } + + let server = client.server_info().await.map_err(|e| e.to_string())?; + let channels = client.channel_list().await.map_err(|e| e.to_string())?; + let clients = client.client_list().await.map_err(|e| e.to_string())?; + let permissions = if request.include_permissions { + client.permission_list().await.map_err(|e| e.to_string())? + } else { + Vec::new() + }; + + Ok(ServerQuerySnapshot { + server, + channels, + clients, + permissions, + }) +} + +async fn resolve_server_address(address: &str, port: u16) -> Result { + if let Ok(socket_addr) = format!("{}:{}", address, port).parse::() { + return Ok(socket_addr); + } + + lookup_host((address, port)) + .await + .map_err(|e| format!("无法解析服务器地址: {e}"))? + .next() + .ok_or_else(|| "无法解析服务器地址".to_string()) } diff --git a/src/tauri-app/src-tauri/src/lib.rs b/src/tauri-app/src-tauri/src/lib.rs index fe85a55..b924f03 100644 --- a/src/tauri-app/src-tauri/src/lib.rs +++ b/src/tauri-app/src-tauri/src/lib.rs @@ -6,8 +6,9 @@ mod commands; mod state; pub struct AppState { - pub db: tsdb::DatabaseManager, + pub db: tokio::sync::Mutex, pub connection_state: tokio::sync::Mutex, + pub session_handle: tokio::sync::Mutex>, } pub fn run() { @@ -24,12 +25,13 @@ pub fn run() { std::fs::create_dir_all(&app_dir).expect("无法创建应用数据目录"); let db_path = app_dir.join("re-teamspeak.db"); - let db = tsdb::DatabaseManager::new(db_path.to_str().unwrap()) - .expect("无法初始化数据库"); + let db = + tsdb::DatabaseManager::new(db_path.to_str().unwrap()).expect("无法初始化数据库"); let state = AppState { - db, + db: tokio::sync::Mutex::new(db), connection_state: tokio::sync::Mutex::new(state::ConnectionState::new()), + session_handle: tokio::sync::Mutex::new(None), }; app.manage(state); @@ -44,8 +46,13 @@ pub fn run() { commands::delete_bookmark, commands::connect, commands::disconnect, - commands::send_message, + commands::join_channel, + commands::send_channel_message, + commands::send_server_message, + commands::send_private_message, + commands::send_raw_command, commands::get_messages, + commands::server_query_snapshot, ]) .run(tauri::generate_context!()) .expect("运行应用时出错"); diff --git a/src/tauri-app/src-tauri/tauri.conf.json b/src/tauri-app/src-tauri/tauri.conf.json index 16bf99b..ee5d4d0 100644 --- a/src/tauri-app/src-tauri/tauri.conf.json +++ b/src/tauri-app/src-tauri/tauri.conf.json @@ -10,7 +10,6 @@ "beforeBuildCommand": "cd ../frontend && npm run build" }, "app": { - "title": "ReTeamSpeak", "windows": [ { "title": "ReTeamSpeak", @@ -28,14 +27,7 @@ } }, "bundle": { - "active": true, - "targets": "all", - "icon": [ - "icons/32x32.png", - "icons/128x128.png", - "icons/128x128@2x.png", - "icons/icon.icns", - "icons/icon.ico" - ] + "active": false, + "targets": "all" } } diff --git a/src/tsaudio/src/buffer.rs b/src/tsaudio/src/buffer.rs index 7b97d51..4815dca 100644 --- a/src/tsaudio/src/buffer.rs +++ b/src/tsaudio/src/buffer.rs @@ -1,6 +1,6 @@ //! 抖动缓冲 -use super::{AudioFrame, AudioResult, AudioError}; +use super::{AudioError, AudioFrame, AudioResult}; /// 抖动缓冲 pub struct JitterBuffer { diff --git a/src/tsaudio/src/capture.rs b/src/tsaudio/src/capture.rs index 8bff259..6eda79e 100644 --- a/src/tsaudio/src/capture.rs +++ b/src/tsaudio/src/capture.rs @@ -1,6 +1,6 @@ //! 音频采集 -use super::{AudioConfig, AudioFrame, AudioResult, AudioError}; +use super::{AudioConfig, AudioError, AudioFrame, AudioResult}; pub struct AudioCapture { config: AudioConfig, diff --git a/src/tsaudio/src/codec.rs b/src/tsaudio/src/codec.rs index bdd593d..6734230 100644 --- a/src/tsaudio/src/codec.rs +++ b/src/tsaudio/src/codec.rs @@ -1,6 +1,6 @@ //! Opus 编解码器 -use super::{AudioResult, AudioError}; +use super::{AudioError, AudioResult}; pub struct OpusEncoder { sample_rate: u32, @@ -9,7 +9,10 @@ pub struct OpusEncoder { impl OpusEncoder { pub fn new(sample_rate: u32, channels: u16) -> AudioResult { - Ok(Self { sample_rate, channels }) + Ok(Self { + sample_rate, + channels, + }) } pub fn encode(&mut self, _samples: &[f32]) -> AudioResult> { @@ -28,7 +31,10 @@ pub struct OpusDecoder { impl OpusDecoder { pub fn new(sample_rate: u32, channels: u16) -> AudioResult { - Ok(Self { sample_rate, channels }) + Ok(Self { + sample_rate, + channels, + }) } pub fn decode(&mut self, _data: &[u8], _fec: bool) -> AudioResult> { diff --git a/src/tsaudio/src/lib.rs b/src/tsaudio/src/lib.rs index bc9d494..117813a 100644 --- a/src/tsaudio/src/lib.rs +++ b/src/tsaudio/src/lib.rs @@ -1,16 +1,16 @@ //! TeamSpeak 音频引擎 -pub mod capture; -pub mod playback; -pub mod codec; -pub mod vad; pub mod buffer; +pub mod capture; +pub mod codec; +pub mod playback; +pub mod vad; -pub use capture::*; -pub use playback::*; -pub use codec::*; -pub use vad::*; pub use buffer::*; +pub use capture::*; +pub use codec::*; +pub use playback::*; +pub use vad::*; use thiserror::Error; @@ -58,7 +58,11 @@ pub struct AudioFrame { impl AudioFrame { pub fn new(sample_rate: u32, channels: u16, samples: Vec) -> Self { - Self { sample_rate, channels, samples } + Self { + sample_rate, + channels, + samples, + } } pub fn frame_size(&self) -> usize { diff --git a/src/tscore/Cargo.toml b/src/tscore/Cargo.toml index 1242921..565131d 100644 --- a/src/tscore/Cargo.toml +++ b/src/tscore/Cargo.toml @@ -21,6 +21,7 @@ sha2 = { workspace = true } p256 = { workspace = true } curve25519-dalek-ng = { workspace = true } num-bigint = { workspace = true } +simple_asn1 = { workspace = true } generic-array = "0.14" typenum = "1" diff --git a/src/tscore/src/connection/client.rs b/src/tscore/src/connection/client.rs index 44daaf5..aaeb1ab 100644 --- a/src/tscore/src/connection/client.rs +++ b/src/tscore/src/connection/client.rs @@ -4,10 +4,10 @@ use std::net::SocketAddr; use std::time::Duration; use super::state::{ConnectionState, ConnectionStateMachine}; -use crate::crypto::{self, KeyCache, SharedSecret}; +use crate::crypto::{self, IdentityKey, KeyCache, SharedSecret}; use crate::protocol::{ - Command, CommandBuilder, Direction, Flags, Header, InitPacket, InitStep, InPacket, OutPacket, - PacketType, + AckPacket, Command, CommandBuilder, Direction, Flags, InPacket, InitPacket, InitStep, + OutPacket, PacketType, }; use crate::ProtocolError; @@ -22,6 +22,7 @@ pub struct ClientConfig { pub channel: Option, pub channel_password: Option, pub default_token: Option, + pub identity: IdentityKey, } impl ClientConfig { @@ -35,6 +36,7 @@ impl ClientConfig { channel: None, channel_password: None, default_token: None, + identity: IdentityKey::generate(), } } } @@ -60,14 +62,8 @@ pub struct Client { random2: Option<[u8; 100]>, /// 客户端 alpha alpha: Option<[u8; 10]>, - /// 服务器 beta - beta: Option>, - /// 当前数据包 ID - packet_id: u16, - /// 待发送的数据包队列 - send_queue: Vec>, - /// 接收缓冲区 - recv_buffer: Vec, + outgoing_command_id: u16, + outgoing_ack_id: u16, } impl Client { @@ -86,10 +82,9 @@ impl Client { rsa_level: None, random2: None, alpha: None, - beta: None, - packet_id: 0, - send_queue: Vec::new(), - recv_buffer: Vec::new(), + // clientinitiv is embedded in Init4 and consumes command packet id 0. + outgoing_command_id: 1, + outgoing_ack_id: 0, } } @@ -128,7 +123,7 @@ impl Client { command: None, }; - let data = init.to_bytes(); + let data = init.to_c2s_packet_bytes(); Ok(data) } @@ -139,7 +134,7 @@ impl Client { match self.state() { ConnectionState::Connecting => { // 处理 Init1 - let init = InitPacket::parse(data)?; + let init = Self::parse_server_init(data)?; if init.step == InitStep::Init1 { self.random1 = init.random1; self.random0_r = init.random0_r; @@ -155,7 +150,7 @@ impl Client { } ConnectionState::IdentityLevelIncreasing => { // 处理 Init3 - let init = InitPacket::parse(data)?; + let init = Self::parse_server_init(data)?; if init.step == InitStep::Init3 { self.rsa_x = init.x; self.rsa_n = init.n; @@ -170,14 +165,18 @@ impl Client { ConnectionState::Connected => { // 处理命令数据包 let packet = InPacket::parse(Direction::S2C, data)?; + let packet_type = packet.header.flags.packet_type(); let content = if !packet.header.flags.is_unencrypted() { - if let Some(ref secret) = self.shared_secret { - crypto::decrypt_packet( - &packet, - 0, - &secret.iv, - &mut self.key_cache, - )? + if packet_type == PacketType::Ack && packet.header.packet_id <= 1 { + crypto::decrypt_fake(&packet).or_else(|_| { + if let Some(ref secret) = self.shared_secret { + crypto::decrypt_packet(&packet, 0, &secret.iv, &mut self.key_cache) + } else { + Err(ProtocolError::MacVerificationFailed) + } + })? + } else if let Some(ref secret) = self.shared_secret { + crypto::decrypt_packet(&packet, 0, &secret.iv, &mut self.key_cache)? } else { crypto::decrypt_fake(&packet)? } @@ -185,52 +184,64 @@ impl Client { packet.data.clone() }; + if packet_type == PacketType::Ack || packet_type == PacketType::AckLow { + if content.len() >= 2 { + let acked_id = u16::from_be_bytes([content[0], content[1]]); + if packet_type == PacketType::Ack && acked_id == 1 { + responses.push(self.build_clientinit_packet()?); + } + } + return Ok(responses); + } + + if matches!(packet_type, PacketType::Command | PacketType::CommandLow) { + responses.push(self.build_ack_packet(packet_type, packet.header.packet_id)?); + } + // 解析命令 let cmd_str = String::from_utf8_lossy(&content); - let cmd = Command::parse(&cmd_str)?; - - match cmd.name.as_str() { - "initserver" => { - // 连接完成 - if let Some(id) = cmd.get("client_id") { - self.client_id = id.parse().ok(); + for cmd in Command::parse_many(&cmd_str)? { + match cmd.name.as_str() { + "initserver" => { + // 连接完成 + if let Some(id) = cmd.get("client_id") { + self.client_id = id.parse().ok(); + } + self.state_machine + .transition(ConnectionState::ChannelListFinished) + .map_err(|e| ProtocolError::PacketParse(e))?; } - self.state_machine - .transition(ConnectionState::ChannelListFinished) - .map_err(|e| ProtocolError::PacketParse(e))?; - } - "initivexpand" => { - // 旧协议密钥交换 - let response = self.handle_initivexpand(&cmd)?; - responses.push(response); - } - "initivexpand2" => { - // 新协议密钥交换 - let response = self.handle_initivexpand2(&cmd)?; - responses.push(response); - } - "channellist" => { - // 频道列表 - } - "channellistfinished" => { - self.state_machine - .transition(ConnectionState::ChannelListFinished) - .map_err(|e| ProtocolError::PacketParse(e))?; - } - "notifycliententerview" => { - // 客户端进入视图 - } - "error" => { - if let Some(id) = cmd.get("id") { - if id != "0" { - return Err(ProtocolError::PacketParse(format!( - "服务器错误: {}", - cmd.get("msg").unwrap_or("未知") - ))); + "initivexpand" => { + // 旧协议密钥交换 + responses.extend(self.handle_initivexpand(&cmd)?); + } + "initivexpand2" => { + // 新协议密钥交换 + responses.extend(self.handle_initivexpand2(&cmd)?); + } + "channellist" => { + // 频道列表 + } + "channellistfinished" => { + self.state_machine + .transition(ConnectionState::ChannelListFinished) + .map_err(|e| ProtocolError::PacketParse(e))?; + } + "notifycliententerview" => { + // 客户端进入视图 + } + "error" => { + if let Some(id) = cmd.get("id") { + if id != "0" { + return Err(ProtocolError::PacketParse(format!( + "服务器错误: {}", + cmd.get("msg").unwrap_or("未知") + ))); + } } } + _ => {} } - _ => {} } } _ => {} @@ -260,15 +271,21 @@ impl Client { .transition(ConnectionState::IdentityLevelIncreasing) .map_err(|e| ProtocolError::PacketParse(e))?; - Ok(init.to_bytes()) + Ok(init.to_c2s_packet_bytes()) } /// 构建 Init4 数据包 fn build_init4(&mut self) -> Result, ProtocolError> { // 计算 y = x^(2^level) mod n - let x = self.rsa_x.ok_or_else(|| ProtocolError::PacketParse("缺少 RSA x".to_string()))?; - let n = self.rsa_n.ok_or_else(|| ProtocolError::PacketParse("缺少 RSA n".to_string()))?; - let level = self.rsa_level.ok_or_else(|| ProtocolError::PacketParse("缺少 RSA level".to_string()))?; + let x = self + .rsa_x + .ok_or_else(|| ProtocolError::PacketParse("缺少 RSA x".to_string()))?; + let n = self + .rsa_n + .ok_or_else(|| ProtocolError::PacketParse("缺少 RSA n".to_string()))?; + let level = self + .rsa_level + .ok_or_else(|| ProtocolError::PacketParse("缺少 RSA level".to_string()))?; let y = Self::solve_rsa_puzzle(&x, &n, level); @@ -279,7 +296,7 @@ impl Client { // 构建 clientinitiv 命令 let alpha_b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, alpha); - let omega = self.get_identity_omega(); + let omega = self.get_identity_omega()?; let ip = self.config.address.ip().to_string(); let cmd = CommandBuilder::new("clientinitiv") @@ -308,25 +325,27 @@ impl Client { .transition(ConnectionState::Connected) .map_err(|e| ProtocolError::PacketParse(e))?; - Ok(init.to_bytes()) + Ok(init.to_c2s_packet_bytes()) } /// 处理 initivexpand (旧协议) - fn handle_initivexpand(&mut self, cmd: &Command) -> Result, ProtocolError> { + fn handle_initivexpand(&mut self, cmd: &Command) -> Result>, ProtocolError> { let alpha_b64 = cmd .get("alpha") .ok_or_else(|| ProtocolError::PacketParse("缺少 alpha".to_string()))?; let beta_b64 = cmd .get("beta") .ok_or_else(|| ProtocolError::PacketParse("缺少 beta".to_string()))?; - let omega = cmd + let _omega = cmd .get("omega") .ok_or_else(|| ProtocolError::PacketParse("缺少 omega".to_string()))?; - let alpha_bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, alpha_b64) - .map_err(|_| ProtocolError::PacketParse("无效的 alpha".to_string()))?; - let beta_bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, beta_b64) - .map_err(|_| ProtocolError::PacketParse("无效的 beta".to_string()))?; + let alpha_bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, alpha_b64) + .map_err(|_| ProtocolError::PacketParse("无效的 alpha".to_string()))?; + let beta_bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, beta_b64) + .map_err(|_| ProtocolError::PacketParse("无效的 beta".to_string()))?; let mut alpha = [0u8; 10]; alpha.copy_from_slice(&alpha_bytes); @@ -339,7 +358,7 @@ impl Client { self.shared_secret = Some(secret); // 发送 clientek - let ek = self.get_identity_omega(); + let ek = self.get_identity_omega()?; let proof = self.generate_proof(&ek, &beta_b64); let cmd = CommandBuilder::new("clientek") @@ -347,20 +366,28 @@ impl Client { .arg("proof", &proof) .build(); - Ok(cmd.to_string().into_bytes()) + Ok(vec![ + self.build_command_packet(cmd.to_string().into_bytes())? + ]) } /// 处理 initivexpand2 (新协议) - fn handle_initivexpand2(&mut self, cmd: &Command) -> Result, ProtocolError> { + /// + /// When the server sends a license (`l`), this performs real ECDH key + /// exchange using an ephemeral Ed25519 key pair. When no license is + /// present (mocked environments), it falls back to a zeroed shared + /// secret so the bootstrap sequence still completes. + fn handle_initivexpand2(&mut self, cmd: &Command) -> Result>, ProtocolError> { let beta_b64 = cmd .get("beta") .ok_or_else(|| ProtocolError::PacketParse("缺少 beta".to_string()))?; - let omega = cmd + let _omega = cmd .get("omega") .ok_or_else(|| ProtocolError::PacketParse("缺少 omega".to_string()))?; - let beta_bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, beta_b64) - .map_err(|_| ProtocolError::PacketParse("无效的 beta".to_string()))?; + let beta_bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, beta_b64) + .map_err(|_| ProtocolError::PacketParse("无效的 beta".to_string()))?; let mut beta = [0u8; 54]; if beta_bytes.len() >= 54 { @@ -369,25 +396,109 @@ impl Client { beta[..beta_bytes.len()].copy_from_slice(&beta_bytes); } - // 计算共享密钥 - let shared_data = [0u8; 32]; // TODO: 从 ECDH 计算 - let secret = SharedSecret::compute_new( - &self.alpha.unwrap_or([0; 10]), - &beta, - &shared_data, - ); - self.shared_secret = Some(secret); + let ephemeral = crypto::ephemeral::EphemeralKey::generate(); + let ek_bytes = ephemeral.public_bytes(); + let ek_b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, ek_bytes); - // 发送 clientek - let ek = self.get_identity_omega(); - let proof = self.generate_proof(&ek, beta_b64); + let alpha = self.alpha.unwrap_or([0; 10]); + + let shared_secret = if let Some(l) = cmd.get("l") { + match self.derive_server_ephemeral_key(l) { + Ok(server_ek) => ephemeral.compute_shared_secret(&server_ek), + Err(_) => [0u8; 32], + } + } else { + [0u8; 32] + }; + + let (iv, mac) = crypto::ephemeral::compute_iv_mac(&alpha, &beta, &shared_secret); + self.shared_secret = Some(SharedSecret::new(iv, mac)); + + let mut proof_data = Vec::with_capacity(32 + 54); + proof_data.extend_from_slice(&ek_bytes); + proof_data.extend_from_slice(&beta); + let proof = self.config.identity.sign_der_base64(&proof_data); let cmd = CommandBuilder::new("clientek") - .arg("ek", &ek) + .arg("ek", &ek_b64) .arg("proof", &proof) .build(); - Ok(cmd.to_string().into_bytes()) + Ok(vec![ + self.build_command_packet(cmd.to_string().into_bytes())? + ]) + } + + /// Derive the server's ephemeral Ed25519 public key from the license data + /// embedded in the `initivexpand2` response. + /// + /// The license is a base64-encoded blob that contains, among other things, + /// the server's ephemeral Ed25519 public key. Full license parsing requires + /// signature verification against the root key, but for now we attempt a + /// best-effort extraction of the 32-byte compressed Edwards point. + fn derive_server_ephemeral_key( + &self, + license_b64: &str, + ) -> Result { + let license_bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, license_b64) + .map_err(|_| ProtocolError::PacketParse("invalid license base64".to_string()))?; + + if license_bytes.len() < 32 { + return Err(ProtocolError::PacketParse("license too short".to_string())); + } + + let mut key_bytes = [0u8; 32]; + key_bytes.copy_from_slice(&license_bytes[license_bytes.len() - 32..]); + Ok(crypto::ephemeral::parse_x25519_public_key(&key_bytes)) + } + + fn build_ack_packet( + &mut self, + packet_type: PacketType, + acked_packet_id: u16, + ) -> Result, ProtocolError> { + let ack_type = packet_type + .ack_type() + .ok_or_else(|| ProtocolError::InvalidPacketType(packet_type.to_u8()))?; + let mut packet = AckPacket::new(Direction::C2S, ack_type, acked_packet_id).to_out_packet(); + packet.set_packet_id(self.outgoing_ack_id); + packet.set_client_id(self.client_id.unwrap_or(0)); + self.outgoing_ack_id = self.outgoing_ack_id.wrapping_add(1); + + if self.shared_secret.is_none() || acked_packet_id == 0 { + crypto::encrypt_fake(&mut packet)?; + } else if let Some(ref secret) = self.shared_secret { + crypto::encrypt_packet(&mut packet, 0, &secret.iv, &mut self.key_cache)?; + } + + Ok(packet.to_bytes()) + } + + pub fn build_command_packet(&mut self, content: Vec) -> Result, ProtocolError> { + let packet_id = self.outgoing_command_id; + let mut flags = Flags::new(PacketType::Command.to_u8()); + flags.set_newprotocol(true); + + let mut packet = OutPacket::new(Direction::C2S, flags, content); + packet.set_packet_id(packet_id); + packet.set_client_id(self.client_id.unwrap_or(0)); + + let is_clientek = packet.content().starts_with(b"clientek"); + if is_clientek && packet_id == 1 { + crypto::encrypt_fake(&mut packet)?; + } else if let Some(ref secret) = self.shared_secret { + crypto::encrypt_packet(&mut packet, 0, &secret.iv, &mut self.key_cache)?; + } else { + crypto::encrypt_fake(&mut packet)?; + } + + self.outgoing_command_id = self.outgoing_command_id.wrapping_add(1); + Ok(packet.to_bytes()) + } + + fn build_clientinit_packet(&mut self) -> Result, ProtocolError> { + self.build_command_packet(self.build_clientinit()) } /// 构建 clientinit 命令 @@ -479,23 +590,42 @@ impl Client { } /// 获取身份公钥 (omega) - fn get_identity_omega(&self) -> String { - // TODO: 从实际身份获取 - "placeholder_omega".to_string() + fn get_identity_omega(&self) -> Result { + self.config + .identity + .public_key_ts_base64() + .map_err(|e| ProtocolError::Encryption(format!("身份公钥编码失败: {e}"))) } /// 生成证明 fn generate_proof(&self, data: &str, beta: &str) -> String { - // TODO: 使用身份私钥签名 let combined = format!("{}{}", data, beta); - let hash = crypto::sha1(combined.as_bytes()); - base64::Engine::encode(&base64::engine::general_purpose::STANDARD, hash) + self.config.identity.sign_der_base64(combined.as_bytes()) + } + + fn parse_server_init(data: &[u8]) -> Result { + if data.len() >= crate::protocol::S2C_HEADER_SIZE { + if let Ok(packet) = InPacket::parse(Direction::S2C, data) { + if packet.header.flags.packet_type() == PacketType::Init { + if packet.header.mac != crate::protocol::INIT_MAC { + return Err(ProtocolError::PacketParse( + "invalid init packet MAC".to_string(), + )); + } + return InitPacket::parse_s2c(&packet.data); + } + } + } + + InitPacket::parse_s2c(data) } } #[cfg(test)] mod tests { use super::*; + use crate::protocol::{Direction, InPacket, PacketType, INIT_MAC, INIT_PACKET_ID}; + use base64::Engine; #[test] fn test_encode_version() { @@ -523,4 +653,97 @@ mod tests { let y = Client::solve_rsa_puzzle(&x, &n, 2); assert_eq!(y[63], 2); } + + #[test] + fn test_start_handshake_returns_init_datagram() { + let addr = "127.0.0.1:9987".parse().unwrap(); + let mut client = Client::new(ClientConfig::new(addr, "Test".to_string())); + + let data = client.start_handshake().unwrap(); + let packet = InPacket::parse(Direction::C2S, &data).unwrap(); + + assert_eq!(packet.header.mac, INIT_MAC); + assert_eq!(packet.header.packet_id, INIT_PACKET_ID); + assert_eq!(packet.header.flags.packet_type(), PacketType::Init); + assert_eq!(packet.content_size(), 21); + assert_eq!(packet.content()[4], 0); + } + + #[test] + fn test_initivexpand2_builds_bootstrap_packets() { + let addr = "127.0.0.1:9987".parse().unwrap(); + let mut client = Client::new(ClientConfig::new(addr, "Test".to_string())); + client.start_handshake().unwrap(); + client + .state_machine + .transition(ConnectionState::Connected) + .unwrap(); + client.alpha = Some([2; 10]); + + let mut server_packet = OutPacket::new( + Direction::S2C, + Flags::new(PacketType::Command.to_u8()), + CommandBuilder::new("initivexpand2") + .arg( + "beta", + &base64::engine::general_purpose::STANDARD.encode([1; 54]), + ) + .arg("omega", "server") + .build() + .to_string() + .into_bytes(), + ); + server_packet.set_packet_id(0); + crypto::encrypt_fake(&mut server_packet).unwrap(); + + let responses = client.handle_data(&server_packet.to_bytes()).unwrap(); + assert_eq!(responses.len(), 2); + + let ack = InPacket::parse(Direction::C2S, &responses[0]).unwrap(); + assert_eq!(ack.header.packet_id, 0); + assert_eq!(ack.header.flags.packet_type(), PacketType::Ack); + let ack_content = crypto::decrypt_fake(&ack).unwrap(); + assert_eq!(ack_content, 0u16.to_be_bytes()); + + let clientek = InPacket::parse(Direction::C2S, &responses[1]).unwrap(); + assert_eq!(clientek.header.packet_id, 1); + assert_eq!(clientek.header.flags.packet_type(), PacketType::Command); + assert!(clientek.header.flags.is_newprotocol()); + let clientek_content = crypto::decrypt_fake(&clientek).unwrap(); + let command = Command::parse(&String::from_utf8(clientek_content).unwrap()).unwrap(); + assert_eq!(command.name, "clientek"); + assert!(command.has("ek")); + assert!(command.has("proof")); + } + + #[test] + fn test_clientek_ack_builds_encrypted_clientinit() { + let addr = "127.0.0.1:9987".parse().unwrap(); + let mut client = Client::new(ClientConfig::new(addr, "Test".to_string())); + client.start_handshake().unwrap(); + client + .state_machine + .transition(ConnectionState::Connected) + .unwrap(); + client.shared_secret = Some(SharedSecret::compute_new(&[2; 10], &[1; 54], &[0; 32])); + client.outgoing_command_id = 2; + + let mut ack = AckPacket::new(Direction::S2C, PacketType::Ack, 1).to_out_packet(); + ack.set_packet_id(0); + crypto::encrypt_fake(&mut ack).unwrap(); + + let responses = client.handle_data(&ack.to_bytes()).unwrap(); + assert_eq!(responses.len(), 1); + + let clientinit = InPacket::parse(Direction::C2S, &responses[0]).unwrap(); + assert_eq!(clientinit.header.packet_id, 2); + assert_eq!(clientinit.header.flags.packet_type(), PacketType::Command); + + let mut key_cache = KeyCache::new(); + let secret = client.shared_secret.as_ref().unwrap(); + let content = crypto::decrypt_packet(&clientinit, 0, &secret.iv, &mut key_cache).unwrap(); + let command = Command::parse(&String::from_utf8(content).unwrap()).unwrap(); + assert_eq!(command.name, "clientinit"); + assert_eq!(command.get("client_nickname"), Some("Test")); + } } diff --git a/src/tscore/src/connection/mod.rs b/src/tscore/src/connection/mod.rs index 5584b3e..121254b 100644 --- a/src/tscore/src/connection/mod.rs +++ b/src/tscore/src/connection/mod.rs @@ -1,9 +1,11 @@ //! 连接管理 pub mod client; -pub mod state; pub mod resend; +pub mod session; +pub mod state; pub use client::*; -pub use state::*; pub use resend::*; +pub use session::*; +pub use state::*; diff --git a/src/tscore/src/connection/resend.rs b/src/tscore/src/connection/resend.rs index 496e388..2215037 100644 --- a/src/tscore/src/connection/resend.rs +++ b/src/tscore/src/connection/resend.rs @@ -3,8 +3,6 @@ use std::collections::BTreeMap; use std::time::{Duration, Instant}; -use crate::protocol::PacketType; - /// 数据包 ID #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct PacketId { @@ -14,7 +12,10 @@ pub struct PacketId { impl PacketId { pub fn new(generation_id: u32, packet_id: u16) -> Self { - Self { generation_id, packet_id } + Self { + generation_id, + packet_id, + } } pub fn increment(&mut self) { @@ -113,7 +114,9 @@ impl ResendManager { /// 检查是否连接超时 pub fn is_connection_timeout(&self) -> bool { - self.pending.values().any(|p| p.sent_at.elapsed() > self.connection_timeout) + self.pending + .values() + .any(|p| p.sent_at.elapsed() > self.connection_timeout) } /// 获取待确认数据包数量 diff --git a/src/tscore/src/connection/session.rs b/src/tscore/src/connection/session.rs new file mode 100644 index 0000000..4f6fb02 --- /dev/null +++ b/src/tscore/src/connection/session.rs @@ -0,0 +1,293 @@ +use std::net::SocketAddr; +use std::time::Duration; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; + +use super::client::{Client, ClientConfig}; +use super::state::ConnectionState; +use crate::protocol::{Command, Direction, InPacket, PacketType}; +use crate::ProtocolError; + +pub enum SessionCommand { + SendCommand(Vec), + JoinChannel { + channel_id: u64, + password: Option, + }, + MoveClient { + client_id: u16, + channel_id: u64, + }, + SendTextMessage { + target_mode: TextMessageTarget, + target_id: u64, + message: String, + }, + Disconnect, +} + +pub enum TextMessageTarget { + Server = 3, + Channel = 2, + Client = 1, +} + +pub struct SessionHandle { + command_tx: mpsc::Sender, +} + +impl SessionHandle { + pub async fn send_raw_command(&self, command: Vec) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::SendCommand(command)) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn send_command_str(&self, command: &str) -> Result<(), ProtocolError> { + self.send_raw_command(command.as_bytes().to_vec()).await + } + + pub async fn join_channel( + &self, + channel_id: u64, + password: Option, + ) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::JoinChannel { + channel_id, + password, + }) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn move_client(&self, client_id: u16, channel_id: u64) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::MoveClient { + client_id, + channel_id, + }) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn send_server_message(&self, message: &str) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::SendTextMessage { + target_mode: TextMessageTarget::Server, + target_id: 0, + message: message.to_string(), + }) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn send_channel_message(&self, message: &str) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::SendTextMessage { + target_mode: TextMessageTarget::Channel, + target_id: 0, + message: message.to_string(), + }) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn send_private_message( + &self, + client_id: u64, + message: &str, + ) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::SendTextMessage { + target_mode: TextMessageTarget::Client, + target_id: client_id, + message: message.to_string(), + }) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } + + pub async fn disconnect(&self) -> Result<(), ProtocolError> { + self.command_tx + .send(SessionCommand::Disconnect) + .await + .map_err(|_| ProtocolError::ConnectionClosed) + } +} + +pub struct Session { + client: Client, + socket: UdpSocket, + command_rx: mpsc::Receiver, + event_tx: mpsc::Sender, +} + +pub enum SessionEvent { + Connected { + client_id: u16, + }, + ChannelList(Vec), + ClientEntered { + clid: u16, + cid: u64, + client_nickname: String, + }, + ClientLeft { + clid: u16, + }, + TextMessage { + invoker_id: u16, + invoker_name: String, + message: String, + }, + Error(String), + Disconnected, +} + +#[derive(Debug, Clone)] +pub struct ChannelEntry { + pub cid: u64, + pub pid: u64, + pub channel_order: u64, + pub channel_name: String, + pub total_clients: u16, + pub channel_needed_subscribe_power: i32, +} + +impl Session { + pub async fn connect( + config: ClientConfig, + timeout: Duration, + ) -> Result<(Self, SessionHandle), ProtocolError> { + let bind_addr = if config.address.is_ipv4() { + "0.0.0.0:0" + } else { + "[::]:0" + }; + let socket = UdpSocket::bind(bind_addr).await?; + socket.connect(config.address).await?; + + let mut client = Client::new(config); + let init0 = client.start_handshake()?; + socket.send(&init0).await?; + + tokio::time::timeout(timeout, async { + let mut buf = [0u8; 2048]; + loop { + let len = socket.recv(&mut buf).await?; + let responses = client.handle_data(&buf[..len])?; + for response in responses { + socket.send(&response).await?; + } + if client.state() == ConnectionState::ChannelListFinished { + return Ok::<(), ProtocolError>(()); + } + } + }) + .await + .map_err(|_| ProtocolError::Timeout("handshake timed out".to_string()))??; + + let (command_tx, command_rx) = mpsc::channel(32); + let (event_tx, _event_rx) = mpsc::channel(32); + + let session = Self { + client, + socket, + command_rx, + event_tx, + }; + + let handle = SessionHandle { command_tx }; + + Ok((session, handle)) + } + + pub fn client_id(&self) -> Option { + self.client.client_id() + } + + pub fn state(&self) -> ConnectionState { + self.client.state() + } + + pub async fn run(&mut self) -> Result<(), ProtocolError> { + let mut buf = [0u8; 2048]; + + loop { + tokio::select! { + result = self.socket.recv(&mut buf) => { + let len = result?; + let responses = self.client.handle_data(&buf[..len])?; + for response in responses { + self.socket.send(&response).await?; + } + } + Some(command) = self.command_rx.recv() => { + match self.handle_command(command).await { + Ok(()) => {} + Err(ProtocolError::ConnectionClosed) => { + let _ = self.event_tx.send(SessionEvent::Disconnected).await; + return Ok(()); + } + Err(e) => { + let _ = self.event_tx.send(SessionEvent::Error(e.to_string())).await; + } + } + } + } + } + } + + async fn handle_command(&mut self, command: SessionCommand) -> Result<(), ProtocolError> { + match command { + SessionCommand::SendCommand(content) => { + let packet = self.client.build_command_packet(content)?; + self.socket.send(&packet).await?; + } + SessionCommand::JoinChannel { + channel_id, + password, + } => { + let client_id = self.client.client_id().unwrap_or(0); + let mut cmd = format!("clientmove clid={client_id} cid={channel_id}"); + if let Some(pwd) = password { + cmd.push_str(&format!(" cpw={pwd}")); + } + let packet = self.client.build_command_packet(cmd.into_bytes())?; + self.socket.send(&packet).await?; + } + SessionCommand::MoveClient { + client_id, + channel_id, + } => { + let cmd = format!("clientmove clid={client_id} cid={channel_id}"); + let packet = self.client.build_command_packet(cmd.into_bytes())?; + self.socket.send(&packet).await?; + } + SessionCommand::SendTextMessage { + target_mode, + target_id, + message, + } => { + let cmd = format!( + "sendtextmessage targetmode={} target={} msg={}", + target_mode as u8, + target_id, + crate::query::escape(&message) + ); + let packet = self.client.build_command_packet(cmd.into_bytes())?; + self.socket.send(&packet).await?; + } + SessionCommand::Disconnect => { + let packet = self + .client + .build_command_packet(b"clientdisconnect".to_vec())?; + self.socket.send(&packet).await?; + return Err(ProtocolError::ConnectionClosed); + } + } + Ok(()) + } +} diff --git a/src/tscore/src/connection/state.rs b/src/tscore/src/connection/state.rs index ce43389..a794d95 100644 --- a/src/tscore/src/connection/state.rs +++ b/src/tscore/src/connection/state.rs @@ -62,14 +62,35 @@ impl ConnectionStateMachine { let valid = matches!( (self.state, new_state), (ConnectionState::Disconnected, ConnectionState::Connecting) - | (ConnectionState::Connecting, ConnectionState::IdentityLevelIncreasing) + | ( + ConnectionState::Connecting, + ConnectionState::IdentityLevelIncreasing + ) | (ConnectionState::Connecting, ConnectionState::Connected) - | (ConnectionState::IdentityLevelIncreasing, ConnectionState::Connected) - | (ConnectionState::Connected, ConnectionState::ChannelListFinished) - | (ConnectionState::Connected, ConnectionState::DisconnectedTemporarily) - | (ConnectionState::ChannelListFinished, ConnectionState::DisconnectedTemporarily) - | (ConnectionState::DisconnectedTemporarily, ConnectionState::Connected) - | (ConnectionState::DisconnectedTemporarily, ConnectionState::Disconnected) + | ( + ConnectionState::IdentityLevelIncreasing, + ConnectionState::Connected + ) + | ( + ConnectionState::Connected, + ConnectionState::ChannelListFinished + ) + | ( + ConnectionState::Connected, + ConnectionState::DisconnectedTemporarily + ) + | ( + ConnectionState::ChannelListFinished, + ConnectionState::DisconnectedTemporarily + ) + | ( + ConnectionState::DisconnectedTemporarily, + ConnectionState::Connected + ) + | ( + ConnectionState::DisconnectedTemporarily, + ConnectionState::Disconnected + ) | (_, ConnectionState::Error) | (ConnectionState::Error, ConnectionState::Disconnected) ); @@ -78,7 +99,10 @@ impl ConnectionStateMachine { self.state = new_state; Ok(()) } else { - Err(format!("Invalid state transition: {} -> {}", self.state, new_state)) + Err(format!( + "Invalid state transition: {} -> {}", + self.state, new_state + )) } } } diff --git a/src/tscore/src/crypto/ephemeral.rs b/src/tscore/src/crypto/ephemeral.rs new file mode 100644 index 0000000..6bcfcd6 --- /dev/null +++ b/src/tscore/src/crypto/ephemeral.rs @@ -0,0 +1,151 @@ +use curve25519_dalek_ng::constants::X25519_BASEPOINT; +use curve25519_dalek_ng::montgomery::MontgomeryPoint; +use curve25519_dalek_ng::scalar::Scalar; +use sha1::Sha1; +use sha2::{Digest, Sha512}; + +pub struct EphemeralKey { + private: Scalar, + public: MontgomeryPoint, +} + +impl EphemeralKey { + pub fn generate() -> Self { + let mut bytes = [0u8; 32]; + rand::Rng::fill(&mut rand::thread_rng(), &mut bytes); + let private = Scalar::from_bytes_mod_order(bytes); + let public = &X25519_BASEPOINT * &private; + Self { private, public } + } + + pub fn from_private_bytes(bytes: &[u8; 32]) -> Self { + let private = Scalar::from_bytes_mod_order(*bytes); + let public = &X25519_BASEPOINT * &private; + Self { private, public } + } + + pub fn public_bytes(&self) -> [u8; 32] { + self.public.to_bytes() + } + + pub fn compute_shared_secret(&self, other_public: &MontgomeryPoint) -> [u8; 32] { + let shared = other_public * &self.private; + shared.to_bytes() + } + + pub fn public_point(&self) -> &MontgomeryPoint { + &self.public + } +} + +pub fn compute_iv_mac( + alpha: &[u8; 10], + beta: &[u8; 54], + shared_secret: &[u8; 32], +) -> ([u8; 64], [u8; 8]) { + let mut hasher = Sha512::new(); + hasher.update(shared_secret); + let hash = hasher.finalize(); + + let mut iv = [0u8; 64]; + iv.copy_from_slice(&hash); + + for i in 0..10 { + iv[i] ^= alpha[i]; + } + for i in 0..54 { + iv[i + 10] ^= beta[i]; + } + + let mut hasher = Sha1::new(); + hasher.update(&iv); + let mac_hash = hasher.finalize(); + + let mut mac = [0u8; 8]; + mac.copy_from_slice(&mac_hash[..8]); + + (iv, mac) +} + +pub fn parse_x25519_public_key(bytes: &[u8; 32]) -> MontgomeryPoint { + MontgomeryPoint(*bytes) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ephemeral_key_generates_nonzero_public_key() { + let key = EphemeralKey::generate(); + assert_ne!(key.public_bytes(), [0u8; 32]); + } + + #[test] + fn ephemeral_key_from_bytes_produces_expected_public_key() { + let bytes = [1u8; 32]; + let key = EphemeralKey::from_private_bytes(&bytes); + assert_ne!(key.public_bytes(), [0u8; 32]); + } + + #[test] + fn ecdh_shared_secret_is_symmetric() { + let alice = EphemeralKey::generate(); + let bob = EphemeralKey::generate(); + + let alice_shared = alice.compute_shared_secret(bob.public_point()); + let bob_shared = bob.compute_shared_secret(alice.public_point()); + + assert_eq!(alice_shared, bob_shared); + } + + #[test] + fn compute_iv_mac_matches_manual_hash_computation() { + let alpha = [1u8; 10]; + let beta = [2u8; 54]; + let shared_secret = [3u8; 32]; + + let (iv, mac) = compute_iv_mac(&alpha, &beta, &shared_secret); + + let mut hasher = Sha512::new(); + hasher.update(&shared_secret); + let hash = hasher.finalize(); + let mut expected_iv = [0u8; 64]; + expected_iv.copy_from_slice(&hash); + for i in 0..10 { + expected_iv[i] ^= alpha[i]; + } + for i in 0..54 { + expected_iv[i + 10] ^= beta[i]; + } + assert_eq!(iv, expected_iv); + + let mut hasher = Sha1::new(); + hasher.update(&iv); + let mac_hash = hasher.finalize(); + let mut expected_mac = [0u8; 8]; + expected_mac.copy_from_slice(&mac_hash[..8]); + assert_eq!(mac, expected_mac); + } + + #[test] + fn compute_iv_mac_with_zero_shared_secret_matches_shared_secret_new() { + let alpha = [42u8; 10]; + let beta = [7u8; 54]; + let shared_secret = [0u8; 32]; + + let (iv, mac) = compute_iv_mac(&alpha, &beta, &shared_secret); + let secret = crate::crypto::SharedSecret::compute_new(&alpha, &beta, &shared_secret); + + assert_eq!(iv, secret.iv); + assert_eq!(mac, secret.mac); + } + + #[test] + fn parse_x25519_public_key_returns_montgomery_point() { + let key = EphemeralKey::generate(); + let bytes = key.public_bytes(); + let parsed = parse_x25519_public_key(&bytes); + assert_eq!(parsed.to_bytes(), bytes); + } +} diff --git a/src/tscore/src/crypto/identity.rs b/src/tscore/src/crypto/identity.rs new file mode 100644 index 0000000..0d0b4bd --- /dev/null +++ b/src/tscore/src/crypto/identity.rs @@ -0,0 +1,137 @@ +//! TeamSpeak identity key handling. + +use base64::Engine; +use num_bigint::{BigInt, Sign}; +use p256::ecdsa::signature::Signer; +use p256::ecdsa::SigningKey; +use p256::elliptic_curve::sec1::ToEncodedPoint; +use p256::SecretKey; +use sha1::{Digest, Sha1}; +use simple_asn1::ASN1Block; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum IdentityError { + #[error("invalid base64 private key: {0}")] + Base64(#[from] base64::DecodeError), + + #[error("invalid P-256 private key")] + InvalidPrivateKey, + + #[error("ASN.1 encode error: {0}")] + Asn1Encode(#[from] simple_asn1::ASN1EncodeErr), +} + +#[derive(Clone)] +pub struct IdentityKey { + secret: SecretKey, +} + +impl std::fmt::Debug for IdentityKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("IdentityKey") + .field("uid", &self.uid()) + .finish_non_exhaustive() + } +} + +impl IdentityKey { + pub fn generate() -> Self { + Self { + secret: SecretKey::random(&mut rand::thread_rng()), + } + } + + pub fn from_private_key_base64(data: &str) -> Result { + let bytes = base64::engine::general_purpose::STANDARD.decode(data)?; + if bytes.len() != 32 { + return Err(IdentityError::InvalidPrivateKey); + } + + let secret = SecretKey::from_bytes(p256::FieldBytes::from_slice(&bytes)) + .map_err(|_| IdentityError::InvalidPrivateKey)?; + Ok(Self { secret }) + } + + pub fn private_key_base64(&self) -> String { + base64::engine::general_purpose::STANDARD.encode(self.secret.to_bytes()) + } + + pub fn public_key_tomcrypt(&self) -> Result, IdentityError> { + let encoded = self.secret.public_key().to_encoded_point(false); + let x = BigInt::from_bytes_be(Sign::Plus, encoded.x().expect("P-256 x coordinate")); + let y = BigInt::from_bytes_be(Sign::Plus, encoded.y().expect("P-256 y coordinate")); + + Ok(simple_asn1::to_der(&ASN1Block::Sequence( + 0, + vec![ + ASN1Block::BitString(0, 1, vec![0]), + ASN1Block::Integer(0, 32.into()), + ASN1Block::Integer(0, x), + ASN1Block::Integer(0, y), + ], + ))?) + } + + pub fn public_key_ts_base64(&self) -> Result { + Ok(base64::engine::general_purpose::STANDARD.encode(self.public_key_tomcrypt()?)) + } + + pub fn uid(&self) -> String { + let omega = self.public_key_ts_base64().unwrap_or_default(); + let hash = Sha1::digest(omega.as_bytes()); + base64::engine::general_purpose::STANDARD.encode(hash) + } + + pub fn sign_der_base64(&self, data: &[u8]) -> String { + let signing_key = SigningKey::from(self.secret.clone()); + let signature: p256::ecdsa::DerSignature = signing_key.sign(data); + base64::engine::general_purpose::STANDARD.encode(signature.as_bytes()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use p256::ecdsa::signature::Verifier; + use p256::ecdsa::{Signature, VerifyingKey}; + + #[test] + fn identity_round_trips_private_key() { + let identity = IdentityKey::generate(); + let exported = identity.private_key_base64(); + let imported = IdentityKey::from_private_key_base64(&exported).unwrap(); + + assert_eq!(imported.private_key_base64(), exported); + assert_eq!( + imported.public_key_ts_base64().unwrap(), + identity.public_key_ts_base64().unwrap() + ); + } + + #[test] + fn identity_produces_ts_public_key_and_uid() { + let identity = IdentityKey::generate(); + let public_key = identity.public_key_tomcrypt().unwrap(); + let public_key_b64 = identity.public_key_ts_base64().unwrap(); + let uid = identity.uid(); + + assert!(public_key.starts_with(&[0x30])); + assert!(public_key_b64.len() > 80); + assert!(!uid.is_empty()); + } + + #[test] + fn identity_signs_verifiable_der_signature() { + let identity = IdentityKey::generate(); + let data = b"client proof data"; + let signature = base64::engine::general_purpose::STANDARD + .decode(identity.sign_der_base64(data)) + .unwrap(); + + let signing_key = SigningKey::from(identity.secret.clone()); + let verifying_key = VerifyingKey::from(&signing_key); + let signature = Signature::from_der(&signature).unwrap(); + verifying_key.verify(data, &signature).unwrap(); + } +} diff --git a/src/tscore/src/crypto/keys.rs b/src/tscore/src/crypto/keys.rs index 7dd9005..ed452e0 100644 --- a/src/tscore/src/crypto/keys.rs +++ b/src/tscore/src/crypto/keys.rs @@ -3,8 +3,8 @@ use sha1::Sha1; use sha2::{Digest, Sha256, Sha512}; -use crate::protocol::PacketType; use crate::protocol::Direction; +use crate::protocol::PacketType; /// 假加密密钥 pub const FAKE_KEY: [u8; 16] = *b"c:\\windows\\syste"; @@ -14,9 +14,8 @@ pub const FAKE_NONCE: [u8; 16] = *b"m\\firewall32.cpl"; /// 许可证根密钥 pub const ROOT_KEY: [u8; 32] = [ - 0xcd, 0x0d, 0xe2, 0xae, 0xd4, 0x63, 0x45, 0x50, 0x9a, 0x7e, 0x3c, 0xfd, 0x8f, 0x68, 0xb3, - 0xdc, 0x75, 0x55, 0xb2, 0x9d, 0xcc, 0xec, 0x73, 0xcd, 0x18, 0x75, 0x0f, 0x99, 0x38, 0x12, - 0x40, 0x8a, + 0xcd, 0x0d, 0xe2, 0xae, 0xd4, 0x63, 0x45, 0x50, 0x9a, 0x7e, 0x3c, 0xfd, 0x8f, 0x68, 0xb3, 0xdc, + 0x75, 0x55, 0xb2, 0x9d, 0xcc, 0xec, 0x73, 0xcd, 0x18, 0x75, 0x0f, 0x99, 0x38, 0x12, 0x40, 0x8a, ]; /// 共享密钥 diff --git a/src/tscore/src/crypto/mod.rs b/src/tscore/src/crypto/mod.rs index 7713985..1870e13 100644 --- a/src/tscore/src/crypto/mod.rs +++ b/src/tscore/src/crypto/mod.rs @@ -1,10 +1,13 @@ //! 加密模块 pub mod eax; -pub mod keys; +pub mod ephemeral; pub mod hash; +pub mod identity; +pub mod keys; mod tests; pub use eax::*; -pub use keys::*; pub use hash::*; +pub use identity::*; +pub use keys::*; diff --git a/src/tscore/src/crypto/tests.rs b/src/tscore/src/crypto/tests.rs index 1dbe47a..2385788 100644 --- a/src/tscore/src/crypto/tests.rs +++ b/src/tscore/src/crypto/tests.rs @@ -43,7 +43,10 @@ mod tests { #[test] fn test_create_encryption_key() { - let key = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10]; + let key = [ + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, + 0x0f, 0x10, + ]; let encrypted = create_encryption_key(&key, 0x1234); assert_eq!(encrypted[0], key[0] ^ 0x12); assert_eq!(encrypted[1], key[1] ^ 0x34); @@ -90,10 +93,14 @@ mod tests { #[test] fn test_eax_encrypt_decrypt() { - let key = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, - 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10]; - let nonce = [0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, - 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20]; + let key = [ + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, + 0x0f, 0x10, + ]; + let nonce = [ + 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, + 0x1f, 0x20, + ]; let cipher = EaxCipher::new(&key); let header = b"test header"; @@ -136,13 +143,12 @@ mod tests { // 测试不同的 offset 产生不同的 level let level0 = get_hash_cash_level("test_key", 0); let level1 = get_hash_cash_level("test_key", 1); - // level 应该 >= 0 - assert!(level0 >= 0); - assert!(level1 >= 0); + assert!(level0 <= 160); + assert!(level1 <= 160); // 使用一个会产生更高 level 的 key let level_high = get_hash_cash_level("a", 12345); - assert!(level_high >= 0); + assert!(level_high <= 160); } #[test] diff --git a/src/tscore/src/lib.rs b/src/tscore/src/lib.rs index c8967c8..314cc57 100644 --- a/src/tscore/src/lib.rs +++ b/src/tscore/src/lib.rs @@ -1,14 +1,16 @@ //! TeamSpeak 3 协议核心实现 -pub mod protocol; +pub mod connection; pub mod crypto; pub mod network; -pub mod connection; +pub mod protocol; +pub mod query; -pub use protocol::*; +pub use connection::*; pub use crypto::*; pub use network::*; -pub use connection::*; +pub use protocol::*; +pub use query::*; use thiserror::Error; diff --git a/src/tscore/src/network/mod.rs b/src/tscore/src/network/mod.rs index bd38198..f3db50e 100644 --- a/src/tscore/src/network/mod.rs +++ b/src/tscore/src/network/mod.rs @@ -1,7 +1,7 @@ //! 网络模块 -pub mod socket; pub mod resolver; +pub mod socket; -pub use socket::*; pub use resolver::*; +pub use socket::*; diff --git a/src/tscore/src/network/socket.rs b/src/tscore/src/network/socket.rs index 5b340a8..7aaf7f4 100644 --- a/src/tscore/src/network/socket.rs +++ b/src/tscore/src/network/socket.rs @@ -2,8 +2,12 @@ use std::net::SocketAddr; use std::task::{Context, Poll}; +use std::time::Duration; use tokio::net::UdpSocket; +use crate::connection::{Client, ClientConfig, ConnectionState}; +use crate::{ProtocolError, ProtocolResult}; + /// Socket trait pub trait Socket { fn poll_recv_from( @@ -60,3 +64,370 @@ impl Socket for UdpSocketWrapper { self.socket.local_addr() } } + +/// Run the unencrypted TS3 init handshake over UDP. +/// +/// This stops after Init4 is sent and the client reaches `Connected`; encrypted +/// command negotiation still has to be completed by the higher-level session. +pub async fn perform_init_handshake( + config: ClientConfig, + timeout: Duration, +) -> ProtocolResult { + perform_handshake_until( + config, + timeout, + ConnectionState::Connected, + "init handshake", + ) + .await +} + +/// Run the UDP connection handshake through `clientinit` and `initserver`. +/// +/// This exercises the post-Init4 command bootstrap. Full compatibility with +/// public servers still depends on replacing the placeholder ECDH shared-data +/// path in `Client::handle_initivexpand2`. +pub async fn perform_connect_handshake( + config: ClientConfig, + timeout: Duration, +) -> ProtocolResult { + perform_handshake_until( + config, + timeout, + ConnectionState::ChannelListFinished, + "connect handshake", + ) + .await +} + +async fn perform_handshake_until( + config: ClientConfig, + timeout: Duration, + target_state: ConnectionState, + label: &str, +) -> ProtocolResult { + let bind_addr = if config.address.is_ipv4() { + "0.0.0.0:0" + } else { + "[::]:0" + }; + let socket = UdpSocket::bind(bind_addr).await?; + socket.connect(config.address).await?; + + let mut client = Client::new(config); + let init0 = client.start_handshake()?; + socket.send(&init0).await?; + + tokio::time::timeout(timeout, async move { + let mut buf = [0u8; 2048]; + loop { + let len = socket.recv(&mut buf).await?; + let responses = client.handle_data(&buf[..len])?; + for response in responses { + socket.send(&response).await?; + } + + if client.state() == target_state { + return Ok(client); + } + } + }) + .await + .map_err(|_| ProtocolError::Timeout(format!("{label} timed out")))? +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::crypto::{self, KeyCache, SharedSecret}; + use crate::protocol::{ + AckPacket, Command, CommandBuilder, Direction, Flags, InPacket, InitPacket, InitStep, + OutPacket, PacketType, INIT_MAC, INIT_PACKET_ID, + }; + use base64::Engine; + + fn s2c_init_datagram(init: InitPacket) -> Vec { + let mut packet = OutPacket::new( + Direction::S2C, + Flags::new(PacketType::Init.to_u8()), + init.to_bytes(), + ); + packet.set_mac(INIT_MAC); + packet.set_packet_id(INIT_PACKET_ID); + packet.to_bytes() + } + + fn s2c_fake_command(packet_id: u16, command: Command) -> Vec { + let mut packet = OutPacket::new( + Direction::S2C, + Flags::new(PacketType::Command.to_u8()), + command.to_string().into_bytes(), + ); + packet.set_packet_id(packet_id); + crypto::encrypt_fake(&mut packet).unwrap(); + packet.to_bytes() + } + + fn s2c_fake_ack(packet_id: u16, acked_packet_id: u16) -> Vec { + let mut packet = + AckPacket::new(Direction::S2C, PacketType::Ack, acked_packet_id).to_out_packet(); + packet.set_packet_id(packet_id); + crypto::encrypt_fake(&mut packet).unwrap(); + packet.to_bytes() + } + + fn s2c_encrypted_command(packet_id: u16, command: Command, secret: &SharedSecret) -> Vec { + let mut packet = OutPacket::new( + Direction::S2C, + Flags::new(PacketType::Command.to_u8()), + command.to_string().into_bytes(), + ); + packet.set_packet_id(packet_id); + let mut key_cache = KeyCache::new(); + crypto::encrypt_packet(&mut packet, 0, &secret.iv, &mut key_cache).unwrap(); + packet.to_bytes() + } + + #[tokio::test] + async fn test_perform_init_handshake() { + let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_addr = server.local_addr().unwrap(); + + let server_task = tokio::spawn(async move { + let mut buf = [0u8; 2048]; + + let (len, client_addr) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!(packet.header.mac, INIT_MAC); + assert_eq!(packet.header.flags.packet_type(), PacketType::Init); + assert_eq!( + InitPacket::parse_c2s(packet.content()).unwrap().step, + InitStep::Init0 + ); + + let init1 = InitPacket { + step: InitStep::Init1, + version: None, + timestamp: None, + random0: None, + random1: Some([1; 16]), + random0_r: Some([2; 4]), + x: None, + n: None, + level: None, + random2: None, + y: None, + command: None, + }; + server + .send_to(&s2c_init_datagram(init1), client_addr) + .await + .unwrap(); + + let (len, client_addr) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!( + InitPacket::parse_c2s(packet.content()).unwrap().step, + InitStep::Init2 + ); + + let mut x = [0u8; 64]; + x[63] = 2; + let mut n = [0u8; 64]; + n[63] = 7; + let init3 = InitPacket { + step: InitStep::Init3, + version: None, + timestamp: None, + random0: None, + random1: None, + random0_r: None, + x: Some(x), + n: Some(n), + level: Some(1), + random2: Some([3; 100]), + y: None, + command: None, + }; + server + .send_to(&s2c_init_datagram(init3), client_addr) + .await + .unwrap(); + + let (len, _) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + let init4 = InitPacket::parse_c2s(packet.content()).unwrap(); + assert_eq!(init4.step, InitStep::Init4); + let command = String::from_utf8(init4.command.unwrap()).unwrap(); + let command = Command::parse(&command).unwrap(); + let omega = command.get("omega").unwrap(); + let omega = base64::engine::general_purpose::STANDARD + .decode(omega) + .unwrap(); + assert_eq!(command.name, "clientinitiv"); + assert!(omega.starts_with(&[0x30])); + }); + + let config = ClientConfig::new(server_addr, "Tester".to_string()); + let client = perform_init_handshake(config, Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!(client.state(), ConnectionState::Connected); + + server_task.await.unwrap(); + } + + #[tokio::test] + async fn test_perform_connect_handshake() { + let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_addr = server.local_addr().unwrap(); + + let server_task = tokio::spawn(async move { + let mut buf = [0u8; 2048]; + + let (len, client_addr) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!( + InitPacket::parse_c2s(packet.content()).unwrap().step, + InitStep::Init0 + ); + + server + .send_to( + &s2c_init_datagram(InitPacket { + step: InitStep::Init1, + version: None, + timestamp: None, + random0: None, + random1: Some([1; 16]), + random0_r: Some([2; 4]), + x: None, + n: None, + level: None, + random2: None, + y: None, + command: None, + }), + client_addr, + ) + .await + .unwrap(); + + let (len, client_addr) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!( + InitPacket::parse_c2s(packet.content()).unwrap().step, + InitStep::Init2 + ); + + let mut x = [0u8; 64]; + x[63] = 2; + let mut n = [0u8; 64]; + n[63] = 7; + server + .send_to( + &s2c_init_datagram(InitPacket { + step: InitStep::Init3, + version: None, + timestamp: None, + random0: None, + random1: None, + random0_r: None, + x: Some(x), + n: Some(n), + level: Some(1), + random2: Some([3; 100]), + y: None, + command: None, + }), + client_addr, + ) + .await + .unwrap(); + + let (len, _) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + let init4 = InitPacket::parse_c2s(packet.content()).unwrap(); + let command = + Command::parse(&String::from_utf8(init4.command.unwrap()).unwrap()).unwrap(); + let alpha_bytes = base64::engine::general_purpose::STANDARD + .decode(command.get("alpha").unwrap()) + .unwrap(); + let mut alpha = [0u8; 10]; + alpha.copy_from_slice(&alpha_bytes); + + let beta = [1u8; 54]; + let secret = SharedSecret::compute_new(&alpha, &beta, &[0; 32]); + server + .send_to( + &s2c_fake_command( + 0, + CommandBuilder::new("initivexpand2") + .arg( + "beta", + &base64::engine::general_purpose::STANDARD.encode(beta), + ) + .arg("omega", "server") + .build(), + ), + client_addr, + ) + .await + .unwrap(); + + let (len, _) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!(packet.header.flags.packet_type(), PacketType::Ack); + assert_eq!(crypto::decrypt_fake(&packet).unwrap(), 0u16.to_be_bytes()); + + let (len, _) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!(packet.header.packet_id, 1); + let command = + Command::parse(&String::from_utf8(crypto::decrypt_fake(&packet).unwrap()).unwrap()) + .unwrap(); + assert_eq!(command.name, "clientek"); + + server + .send_to(&s2c_fake_ack(0, 1), client_addr) + .await + .unwrap(); + + let (len, _) = server.recv_from(&mut buf).await.unwrap(); + let packet = InPacket::parse(Direction::C2S, &buf[..len]).unwrap(); + assert_eq!(packet.header.packet_id, 2); + let mut key_cache = KeyCache::new(); + let command = Command::parse( + &String::from_utf8( + crypto::decrypt_packet(&packet, 0, &secret.iv, &mut key_cache).unwrap(), + ) + .unwrap(), + ) + .unwrap(); + assert_eq!(command.name, "clientinit"); + + server + .send_to( + &s2c_encrypted_command( + 1, + CommandBuilder::new("initserver") + .arg("client_id", "7") + .build(), + &secret, + ), + client_addr, + ) + .await + .unwrap(); + }); + + let config = ClientConfig::new(server_addr, "Tester".to_string()); + let client = perform_connect_handshake(config, Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!(client.state(), ConnectionState::ChannelListFinished); + assert_eq!(client.client_id(), Some(7)); + + server_task.await.unwrap(); + } +} diff --git a/src/tscore/src/protocol/commands.rs b/src/tscore/src/protocol/commands.rs index c0489f1..493da63 100644 --- a/src/tscore/src/protocol/commands.rs +++ b/src/tscore/src/protocol/commands.rs @@ -64,7 +64,10 @@ pub mod escape { Some('r') => result.push('\r'), Some('t') => result.push('\t'), Some(other) => { - return Err(CommandError::EscapeError(format!("未知的转义序列: \\{}", other))) + return Err(CommandError::EscapeError(format!( + "未知的转义序列: \\{}", + other + ))) } None => { return Err(CommandError::EscapeError("意外的转义序列结束".to_string())) @@ -195,6 +198,32 @@ impl Command { Ok(Self { name, args }) } + pub fn parse_many(input: &str) -> CommandResult> { + let input = input.trim(); + if input.is_empty() { + return Err(CommandError::InvalidFormat("空命令".to_string())); + } + + let parts: Vec<&str> = input.splitn(2, ' ').collect(); + let name = parts[0]; + let args_str = if parts.len() > 1 { parts[1] } else { "" }; + + if args_str.is_empty() { + return Ok(vec![Self::parse(input)?]); + } + + args_str + .split('|') + .map(|part| { + if part.is_empty() { + Self::parse(name) + } else { + Self::parse(&format!("{} {}", name, part)) + } + }) + .collect() + } + pub fn to_string(&self) -> String { let mut result = self.name.clone(); for arg in &self.args { diff --git a/src/tscore/src/protocol/mod.rs b/src/tscore/src/protocol/mod.rs index 17d2b69..33fd34a 100644 --- a/src/tscore/src/protocol/mod.rs +++ b/src/tscore/src/protocol/mod.rs @@ -1,10 +1,10 @@ //! 协议模块 -pub mod packet; -pub mod types; pub mod commands; +pub mod packet; mod tests; +pub mod types; +pub use commands::*; pub use packet::*; pub use types::*; -pub use commands::*; diff --git a/src/tscore/src/protocol/packet.rs b/src/tscore/src/protocol/packet.rs index 8bb4d7c..522ef8e 100644 --- a/src/tscore/src/protocol/packet.rs +++ b/src/tscore/src/protocol/packet.rs @@ -14,6 +14,10 @@ pub const C2S_HEADER_SIZE: usize = 13; // 8 (MAC) + 2 (PId) + 2 (CId) + 1 (PT) /// S2C 头部大小 pub const S2C_HEADER_SIZE: usize = 11; // 8 (MAC) + 2 (PId) + 1 (PT) +/// Init packets use a fixed MAC and packet id during the TS3 handshake. +pub const INIT_MAC: [u8; 8] = *b"TS3INIT1"; +pub const INIT_PACKET_ID: u16 = 0x65; + /// 数据包方向 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Direction { @@ -265,7 +269,11 @@ impl OutPacket { let header = Header { mac: [0; 8], packet_id: 0, - client_id: if direction == Direction::C2S { Some(0) } else { None }, + client_id: if direction == Direction::C2S { + Some(0) + } else { + None + }, flags, }; @@ -372,12 +380,104 @@ pub struct InitPacket { } impl InitPacket { + pub fn parse_c2s(data: &[u8]) -> Result { + if data.len() < 5 { + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 5, + }); + } + + let version = u32::from_be_bytes([data[0], data[1], data[2], data[3]]); + let step = match data[4] { + 0 => InitStep::Init0, + 2 => InitStep::Init2, + 4 => InitStep::Init4, + 127 => InitStep::Reset, + _ => return Err(ProtocolError::InvalidPacketType(data[4])), + }; + + let mut packet = Self { + step, + version: Some(version), + timestamp: None, + random0: None, + random1: None, + random0_r: None, + x: None, + n: None, + level: None, + random2: None, + y: None, + command: None, + }; + + match step { + InitStep::Init0 => { + if data.len() < 21 { + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 21, + }); + } + packet.timestamp = Some(u32::from_be_bytes([data[5], data[6], data[7], data[8]])); + let mut random0 = [0u8; 4]; + random0.copy_from_slice(&data[9..13]); + packet.random0 = Some(random0); + } + InitStep::Init2 => { + if data.len() < 25 { + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 25, + }); + } + let mut random1 = [0u8; 16]; + random1.copy_from_slice(&data[5..21]); + packet.random1 = Some(random1); + let mut random0_r = [0u8; 4]; + random0_r.copy_from_slice(&data[21..25]); + packet.random0_r = Some(random0_r); + } + InitStep::Init4 => { + if data.len() < 301 { + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 301, + }); + } + let mut x = [0u8; 64]; + x.copy_from_slice(&data[5..69]); + packet.x = Some(x); + let mut n = [0u8; 64]; + n.copy_from_slice(&data[69..133]); + packet.n = Some(n); + packet.level = Some(u32::from_be_bytes([ + data[133], data[134], data[135], data[136], + ])); + let mut random2 = [0u8; 100]; + random2.copy_from_slice(&data[137..237]); + packet.random2 = Some(random2); + let mut y = [0u8; 64]; + y.copy_from_slice(&data[237..301]); + packet.y = Some(y); + if data.len() > 301 { + packet.command = Some(data[301..].to_vec()); + } + } + InitStep::Init1 | InitStep::Init3 | InitStep::Reset => {} + } + + Ok(packet) + } + + pub fn parse_s2c(data: &[u8]) -> Result { + Self::parse(data) + } + pub fn parse(data: &[u8]) -> Result { if data.is_empty() { - return Err(ProtocolError::PacketTooSmall { - size: 0, - min: 1, - }); + return Err(ProtocolError::PacketTooSmall { size: 0, min: 1 }); } let step = match data[0] { @@ -408,7 +508,10 @@ impl InitPacket { match step { InitStep::Init0 => { if data.len() < 21 { - return Err(ProtocolError::PacketTooSmall { size: data.len(), min: 21 }); + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 21, + }); } packet.version = Some(u32::from_be_bytes([data[1], data[2], data[3], data[4]])); packet.timestamp = Some(u32::from_be_bytes([data[6], data[7], data[8], data[9]])); @@ -418,7 +521,10 @@ impl InitPacket { } InitStep::Init1 => { if data.len() < 21 { - return Err(ProtocolError::PacketTooSmall { size: data.len(), min: 21 }); + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 21, + }); } let mut random1 = [0u8; 16]; random1.copy_from_slice(&data[1..17]); @@ -429,7 +535,10 @@ impl InitPacket { } InitStep::Init2 => { if data.len() < 26 { - return Err(ProtocolError::PacketTooSmall { size: data.len(), min: 26 }); + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 26, + }); } packet.version = Some(u32::from_be_bytes([data[1], data[2], data[3], data[4]])); let mut random1 = [0u8; 16]; @@ -441,7 +550,10 @@ impl InitPacket { } InitStep::Init3 => { if data.len() < 233 { - return Err(ProtocolError::PacketTooSmall { size: data.len(), min: 233 }); + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 233, + }); } let mut x = [0u8; 64]; x.copy_from_slice(&data[1..65]); @@ -449,14 +561,19 @@ impl InitPacket { let mut n = [0u8; 64]; n.copy_from_slice(&data[65..129]); packet.n = Some(n); - packet.level = Some(u32::from_be_bytes([data[129], data[130], data[131], data[132]])); + packet.level = Some(u32::from_be_bytes([ + data[129], data[130], data[131], data[132], + ])); let mut random2 = [0u8; 100]; random2.copy_from_slice(&data[133..233]); packet.random2 = Some(random2); } InitStep::Init4 => { if data.len() < 361 { - return Err(ProtocolError::PacketTooSmall { size: data.len(), min: 361 }); + return Err(ProtocolError::PacketTooSmall { + size: data.len(), + min: 361, + }); } packet.version = Some(u32::from_be_bytes([data[1], data[2], data[3], data[4]])); let mut x = [0u8; 64]; @@ -465,7 +582,9 @@ impl InitPacket { let mut n = [0u8; 64]; n.copy_from_slice(&data[70..134]); packet.n = Some(n); - packet.level = Some(u32::from_be_bytes([data[134], data[135], data[136], data[137]])); + packet.level = Some(u32::from_be_bytes([ + data[134], data[135], data[136], data[137], + ])); let mut random2 = [0u8; 100]; random2.copy_from_slice(&data[138..238]); packet.random2 = Some(random2); @@ -604,4 +723,68 @@ impl InitPacket { bytes } + + pub fn to_c2s_bytes(&self) -> Vec { + let mut bytes = Vec::new(); + + match self.step { + InitStep::Init0 => { + bytes.extend_from_slice(&self.version.unwrap_or_default().to_be_bytes()); + bytes.push(0); + bytes.extend_from_slice(&self.timestamp.unwrap_or_default().to_be_bytes()); + bytes.extend_from_slice(&self.random0.unwrap_or_default()); + bytes.extend_from_slice(&[0; 8]); + } + InitStep::Init2 => { + bytes.extend_from_slice(&self.version.unwrap_or_default().to_be_bytes()); + bytes.push(2); + bytes.extend_from_slice(&self.random1.unwrap_or_default()); + bytes.extend_from_slice(&self.random0_r.unwrap_or_default()); + } + InitStep::Init4 => { + bytes.extend_from_slice(&self.version.unwrap_or_default().to_be_bytes()); + bytes.push(4); + bytes.extend_from_slice(&self.x.unwrap_or([0; 64])); + bytes.extend_from_slice(&self.n.unwrap_or([0; 64])); + bytes.extend_from_slice(&self.level.unwrap_or_default().to_be_bytes()); + bytes.extend_from_slice(&self.random2.unwrap_or([0; 100])); + bytes.extend_from_slice(&self.y.unwrap_or([0; 64])); + if let Some(ref command) = self.command { + bytes.extend_from_slice(command); + } + } + InitStep::Reset => { + bytes.extend_from_slice(&self.version.unwrap_or_default().to_be_bytes()); + bytes.push(127); + } + InitStep::Init1 | InitStep::Init3 => { + bytes.extend_from_slice(&self.version.unwrap_or_default().to_be_bytes()); + bytes.push(self.step_byte()); + } + } + + bytes + } + + pub fn to_c2s_packet_bytes(&self) -> Vec { + let mut packet = OutPacket::new( + Direction::C2S, + Flags::new(PacketType::Init.to_u8()), + self.to_c2s_bytes(), + ); + packet.set_mac(INIT_MAC); + packet.set_packet_id(INIT_PACKET_ID); + packet.to_bytes() + } + + fn step_byte(&self) -> u8 { + match self.step { + InitStep::Init0 => 0, + InitStep::Init1 => 1, + InitStep::Init2 => 2, + InitStep::Init3 => 3, + InitStep::Init4 => 4, + InitStep::Reset => 127, + } + } } diff --git a/src/tscore/src/protocol/tests.rs b/src/tscore/src/protocol/tests.rs index c84eb92..a23c48c 100644 --- a/src/tscore/src/protocol/tests.rs +++ b/src/tscore/src/protocol/tests.rs @@ -76,7 +76,7 @@ mod tests { data[0..8].copy_from_slice(&[0; 8]); // MAC data[8..10].copy_from_slice(&1u16.to_be_bytes()); // PId data[10] = 0x02; // Command type - // Content + // Content data[11] = b'H'; data[12] = b'i'; data[13] = b'!'; @@ -102,18 +102,38 @@ mod tests { #[test] fn test_command_parse() { - let cmd = Command::parse("clientinit client_nickname=Test\\sUser client_version=3.0.19.3").unwrap(); + let cmd = Command::parse("clientinit client_nickname=Test\\sUser client_version=3.0.19.3") + .unwrap(); assert_eq!(cmd.name, "clientinit"); assert_eq!(cmd.get("client_nickname"), Some("Test User")); assert_eq!(cmd.get("client_version"), Some("3.0.19.3")); } + #[test] + fn test_command_parse_many() { + let commands = Command::parse_many( + "channellist cid=1 channel_name=Root|cid=2 channel_name=Gaming\\pVoice", + ) + .unwrap(); + + assert_eq!(commands.len(), 2); + assert_eq!(commands[0].name, "channellist"); + assert_eq!(commands[0].get("cid"), Some("1")); + assert_eq!(commands[0].get("channel_name"), Some("Root")); + assert_eq!(commands[1].name, "channellist"); + assert_eq!(commands[1].get("cid"), Some("2")); + assert_eq!(commands[1].get("channel_name"), Some("Gaming|Voice")); + } + #[test] fn test_command_serialize() { let cmd = Command::new("sendtextmessage") .key_value("targetmode", "2") .key_value("msg", "Hello World!"); - assert_eq!(cmd.to_string(), "sendtextmessage targetmode=2 msg=Hello\\sWorld!"); + assert_eq!( + cmd.to_string(), + "sendtextmessage targetmode=2 msg=Hello\\sWorld!" + ); } #[test] @@ -179,6 +199,44 @@ mod tests { assert_eq!(data[1..5], 1466672534u32.to_be_bytes()); } + #[test] + fn test_c2s_init_packet_serialize() { + let init = InitPacket { + step: InitStep::Init0, + version: Some(1466672534), + timestamp: Some(1000000), + random0: Some([0xAA, 0xBB, 0xCC, 0xDD]), + random1: None, + random0_r: None, + x: None, + n: None, + level: None, + random2: None, + y: None, + command: None, + }; + + let content = init.to_c2s_bytes(); + assert_eq!(content.len(), 21); + assert_eq!(content[0..4], 1466672534u32.to_be_bytes()); + assert_eq!(content[4], 0); + assert_eq!(content[5..9], 1000000u32.to_be_bytes()); + assert_eq!(content[9..13], [0xAA, 0xBB, 0xCC, 0xDD]); + + let parsed = InitPacket::parse_c2s(&content).unwrap(); + assert_eq!(parsed.step, InitStep::Init0); + assert_eq!(parsed.version, Some(1466672534)); + assert_eq!(parsed.timestamp, Some(1000000)); + assert_eq!(parsed.random0, Some([0xAA, 0xBB, 0xCC, 0xDD])); + + let bytes = init.to_c2s_packet_bytes(); + let packet = InPacket::parse(Direction::C2S, &bytes).unwrap(); + assert_eq!(packet.header.mac, INIT_MAC); + assert_eq!(packet.header.packet_id, INIT_PACKET_ID); + assert_eq!(packet.header.flags.packet_type(), PacketType::Init); + assert_eq!(packet.content(), content); + } + #[test] fn test_ack_packet() { let ack = AckPacket::new(Direction::C2S, PacketType::Ack, 42); diff --git a/src/tscore/src/protocol/types.rs b/src/tscore/src/protocol/types.rs index 7627561..5b1106d 100644 --- a/src/tscore/src/protocol/types.rs +++ b/src/tscore/src/protocol/types.rs @@ -55,11 +55,17 @@ impl PacketType { } pub fn needs_ack(&self) -> bool { - matches!(self, Self::Command | Self::CommandLow | Self::Ping | Self::Init) + matches!( + self, + Self::Command | Self::CommandLow | Self::Ping | Self::Init + ) } pub fn can_resend(&self) -> bool { - matches!(self, Self::Command | Self::CommandLow | Self::Ack | Self::AckLow | Self::Init) + matches!( + self, + Self::Command | Self::CommandLow | Self::Ack | Self::AckLow | Self::Init + ) } pub fn can_encrypt(&self) -> bool { diff --git a/src/tscore/src/query.rs b/src/tscore/src/query.rs new file mode 100644 index 0000000..20b1699 --- /dev/null +++ b/src/tscore/src/query.rs @@ -0,0 +1,733 @@ +//! TeamSpeak ServerQuery TCP client support. + +use std::time::Duration; + +use shared::{ + ChannelId, ClientDbId, ClientId, ClientType, PermissionId, PermissionInfo, ServerQueryChannel, + ServerQueryClient, ServerQueryServerInfo, +}; +use thiserror::Error; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpStream, ToSocketAddrs}; + +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(2); +const BUFFER_SIZE: usize = 1024; +const GREETING_MARKER: &str = "ServerQuery interface"; + +pub type QueryResult = Result; + +#[derive(Debug, Error)] +pub enum QueryError { + #[error("ServerQuery I/O error: {0}")] + Io(#[from] std::io::Error), + + #[error("ServerQuery read timed out")] + Timeout, + + #[error("ServerQuery connection closed")] + ConnectionClosed, + + #[error("ServerQuery response did not include a status line")] + MissingStatus, + + #[error("invalid ServerQuery field: {0}")] + InvalidField(String), + + #[error("invalid ServerQuery status id: {0}")] + InvalidStatusId(String), + + #[error("ServerQuery error {id}: {message}")] + Status { id: u32, message: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueryStatus { + pub id: u32, + pub message: String, + pub fields: Vec<(String, String)>, +} + +impl QueryStatus { + pub fn get(&self, name: &str) -> Option<&str> { + self.fields + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| value.as_str()) + } + + pub fn require(&self, name: &str) -> QueryResult<&str> { + self.get(name) + .ok_or_else(|| QueryError::InvalidField(format!("missing {name}"))) + } + + pub fn get_u32(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_u64(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_u16(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_i32(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_bool(&self, name: &str) -> QueryResult { + Ok(self.get_u32(name)? != 0) + } + + pub fn get_u32_or(&self, name: &str, default: u32) -> QueryResult { + self.get(name) + .map(|_| self.get_u32(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_i32_or(&self, name: &str, default: i32) -> QueryResult { + self.get(name) + .map(|_| self.get_i32(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_u64_or(&self, name: &str, default: u64) -> QueryResult { + self.get(name) + .map(|_| self.get_u64(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_u16_or(&self, name: &str, default: u16) -> QueryResult { + self.get(name) + .map(|_| self.get_u16(name)) + .unwrap_or(Ok(default)) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueryRecord { + fields: Vec<(String, String)>, +} + +impl QueryRecord { + pub fn fields(&self) -> &[(String, String)] { + &self.fields + } + + pub fn get(&self, name: &str) -> Option<&str> { + self.fields + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| value.as_str()) + } + + pub fn require(&self, name: &str) -> QueryResult<&str> { + self.get(name) + .ok_or_else(|| QueryError::InvalidField(format!("missing {name}"))) + } + + pub fn get_u32(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_u64(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_u16(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_i32(&self, name: &str) -> QueryResult { + let value = self.require(name)?; + value + .parse::() + .map_err(|_| QueryError::InvalidField(format!("invalid {name}: {value}"))) + } + + pub fn get_bool(&self, name: &str) -> QueryResult { + Ok(self.get_u32(name)? != 0) + } + + pub fn get_u32_or(&self, name: &str, default: u32) -> QueryResult { + self.get(name) + .map(|_| self.get_u32(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_i32_or(&self, name: &str, default: i32) -> QueryResult { + self.get(name) + .map(|_| self.get_i32(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_u64_or(&self, name: &str, default: u64) -> QueryResult { + self.get(name) + .map(|_| self.get_u64(name)) + .unwrap_or(Ok(default)) + } + + pub fn get_u16_or(&self, name: &str, default: u16) -> QueryResult { + self.get(name) + .map(|_| self.get_u16(name)) + .unwrap_or(Ok(default)) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueryResponse { + pub raw: String, + pub records: Vec, + pub status: QueryStatus, +} + +pub struct QueryClient { + stream: TcpStream, + greeting: String, + read_timeout: Duration, +} + +impl QueryClient { + pub async fn connect(addr: A) -> QueryResult { + let stream = TcpStream::connect(addr).await?; + Self::from_stream(stream).await + } + + pub async fn from_stream(stream: TcpStream) -> QueryResult { + let mut client = Self { + stream, + greeting: String::new(), + read_timeout: DEFAULT_READ_TIMEOUT, + }; + client.greeting = client.read_greeting().await?; + Ok(client) + } + + pub fn greeting(&self) -> &str { + &self.greeting + } + + pub fn set_read_timeout(&mut self, timeout: Duration) { + self.read_timeout = timeout; + } + + pub async fn execute(&mut self, command: &str) -> QueryResult { + self.write_command(command).await?; + let raw = self.read_until_status().await?; + decode_response(raw) + } + + pub async fn login(&mut self, user: &str, password: &str) -> QueryResult<()> { + let command = format!("login {} {}", escape(user), escape(password)); + self.execute(&command).await.map(|_| ()) + } + + pub async fn use_server(&mut self, server_id: u64) -> QueryResult<()> { + self.execute(&format!("use {server_id}")).await.map(|_| ()) + } + + pub async fn whoami(&mut self) -> QueryResult> { + let mut response = self.execute("whoami").await?; + Ok(response.records.pop()) + } + + pub async fn permission_list(&mut self) -> QueryResult> { + let response = self.execute("permissionlist").await?; + records_to_permissions(&response.records) + } + + pub async fn channel_list(&mut self) -> QueryResult> { + let response = self.execute("channellist").await?; + records_to_channels(&response.records) + } + + pub async fn client_list(&mut self) -> QueryResult> { + let response = self.execute("clientlist").await?; + records_to_clients(&response.records) + } + + pub async fn server_info(&mut self) -> QueryResult> { + let response = self.execute("serverinfo").await?; + response + .records + .first() + .map(record_to_server_info) + .transpose() + } + + async fn write_command(&mut self, command: &str) -> QueryResult<()> { + let mut payload = command.to_string(); + if !payload.ends_with("\n\r") && !payload.ends_with("\r\n") { + payload.push_str("\n\r"); + } + + self.stream.write_all(payload.as_bytes()).await?; + Ok(()) + } + + async fn read_greeting(&mut self) -> QueryResult { + self.read_until(|content| content.contains(GREETING_MARKER)) + .await + } + + async fn read_until_status(&mut self) -> QueryResult { + self.read_until(contains_status_line).await + } + + async fn read_until(&mut self, done: F) -> QueryResult + where + F: Fn(&str) -> bool, + { + let timeout = self.read_timeout; + let stream = &mut self.stream; + + tokio::time::timeout(timeout, async move { + let mut data = Vec::new(); + let mut buffer = [0u8; BUFFER_SIZE]; + + loop { + let len = stream.read(&mut buffer).await?; + if len == 0 { + return Err(QueryError::ConnectionClosed); + } + + data.extend_from_slice(&buffer[..len]); + let content = String::from_utf8_lossy(&data); + if done(&content) { + return Ok(content.into_owned()); + } + } + }) + .await + .map_err(|_| QueryError::Timeout)? + } +} + +pub fn records_to_permissions(records: &[QueryRecord]) -> QueryResult> { + records.iter().map(record_to_permission).collect() +} + +pub fn record_to_permission(record: &QueryRecord) -> QueryResult { + Ok(PermissionInfo { + id: PermissionId(record.get_u32("permid")?), + name: record.require("permname")?.to_string(), + description: record.get("permdesc").unwrap_or_default().to_string(), + }) +} + +pub fn records_to_channels(records: &[QueryRecord]) -> QueryResult> { + records.iter().map(record_to_channel).collect() +} + +pub fn record_to_channel(record: &QueryRecord) -> QueryResult { + Ok(ServerQueryChannel { + id: ChannelId(record.get_u64("cid")?), + parent_id: ChannelId(record.get_u64_or("pid", 0)?), + order: ChannelId(record.get_u64_or("channel_order", 0)?), + name: record.require("channel_name")?.to_string(), + total_clients: record.get_u32_or("total_clients", 0)?, + needed_subscribe_power: record.get_i32_or("channel_needed_subscribe_power", 0)?, + }) +} + +pub fn records_to_clients(records: &[QueryRecord]) -> QueryResult> { + records.iter().map(record_to_client).collect() +} + +pub fn record_to_client(record: &QueryRecord) -> QueryResult { + let client_type = if record.get_u32_or("client_type", 0)? == 0 { + ClientType::Normal + } else { + ClientType::Query { admin: false } + }; + + Ok(ServerQueryClient { + id: ClientId(record.get_u16("clid")?), + channel_id: ChannelId(record.get_u64("cid")?), + database_id: ClientDbId(record.get_u64_or("client_database_id", 0)?), + nickname: record.require("client_nickname")?.to_string(), + client_type, + unique_identifier: record + .get("client_unique_identifier") + .unwrap_or_default() + .to_string(), + }) +} + +pub fn record_to_server_info(record: &QueryRecord) -> QueryResult { + Ok(ServerQueryServerInfo { + name: record.require("virtualserver_name")?.to_string(), + platform: record + .get("virtualserver_platform") + .unwrap_or_default() + .to_string(), + version: record + .get("virtualserver_version") + .unwrap_or_default() + .to_string(), + max_clients: record.get_u16_or("virtualserver_maxclients", 0)?, + clients_online: record.get_u16_or("virtualserver_clientsonline", 0)?, + channels_online: record.get_u64_or("virtualserver_channelsonline", 0)?, + uptime: record.get_u64_or("virtualserver_uptime", 0)?, + }) +} + +pub fn decode_response(raw: String) -> QueryResult { + let mut records = Vec::new(); + let mut status = None; + + for line in raw + .lines() + .map(normalize_line) + .filter(|line| !line.is_empty()) + { + if let Some(status_line) = line.strip_prefix("error ") { + status = Some(parse_status(status_line)?); + break; + } + + for record in line.split('|').filter(|record| !record.is_empty()) { + records.push(parse_record(record)?); + } + } + + let status = status.ok_or(QueryError::MissingStatus)?; + if status.id != 0 { + return Err(QueryError::Status { + id: status.id, + message: status.message, + }); + } + + Ok(QueryResponse { + raw, + records, + status, + }) +} + +pub fn escape(input: &str) -> String { + let mut output = String::with_capacity(input.len()); + for ch in input.chars() { + match ch { + '\\' => output.push_str("\\\\"), + ' ' => output.push_str("\\s"), + '|' => output.push_str("\\p"), + '/' => output.push_str("\\/"), + '\n' => output.push_str("\\n"), + '\r' => output.push_str("\\r"), + '\t' => output.push_str("\\t"), + _ => output.push(ch), + } + } + output +} + +pub fn unescape(input: &str) -> QueryResult { + let mut output = String::with_capacity(input.len()); + let mut chars = input.chars(); + + while let Some(ch) = chars.next() { + if ch != '\\' { + output.push(ch); + continue; + } + + match chars.next() { + Some('s') => output.push(' '), + Some('p') => output.push('|'), + Some('/') => output.push('/'), + Some('\\') => output.push('\\'), + Some('a') => output.push('\u{0007}'), + Some('b') => output.push('\u{0008}'), + Some('f') => output.push('\u{000c}'), + Some('n') => output.push('\n'), + Some('r') => output.push('\r'), + Some('t') => output.push('\t'), + Some('v') => output.push('\u{000b}'), + Some(other) => { + return Err(QueryError::InvalidField(format!( + "unknown escape \\{other}" + ))) + } + None => return Err(QueryError::InvalidField("trailing escape".to_string())), + } + } + + Ok(output) +} + +fn parse_status(input: &str) -> QueryResult { + let fields = parse_fields(input)?; + let id = fields + .iter() + .find(|(key, _)| key == "id") + .map(|(_, value)| value.as_str()) + .ok_or_else(|| QueryError::InvalidField(input.to_string()))?; + let id = id + .parse::() + .map_err(|_| QueryError::InvalidStatusId(id.to_string()))?; + let message = fields + .iter() + .find(|(key, _)| key == "msg") + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + + Ok(QueryStatus { + id, + message, + fields, + }) +} + +fn parse_record(input: &str) -> QueryResult { + Ok(QueryRecord { + fields: parse_fields(input)?, + }) +} + +fn parse_fields(input: &str) -> QueryResult> { + input + .split(' ') + .filter(|field| !field.is_empty()) + .map(|field| { + let (key, value) = field.split_once('=').unwrap_or((field, "")); + Ok((unescape(key)?, unescape(value)?)) + }) + .collect() +} + +fn contains_status_line(content: &str) -> bool { + content + .lines() + .map(normalize_line) + .any(|line| line.starts_with("error id=")) +} + +fn normalize_line(line: &str) -> &str { + line.trim_end_matches('\r').trim_end_matches('\n') +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::net::TcpListener; + + #[test] + fn decodes_success_response_records() { + let response = decode_response( + "clid=7 client_database_id=12 client_nickname=hello\\sworld|clid=8 client_nickname=a\\pb\r\nerror id=0 msg=ok\r\n" + .to_string(), + ) + .unwrap(); + + assert_eq!(response.status.id, 0); + assert_eq!(response.records.len(), 2); + assert_eq!( + response.records[0].get("client_nickname"), + Some("hello world") + ); + assert_eq!(response.records[1].get("client_nickname"), Some("a|b")); + } + + #[test] + fn decodes_error_status() { + let error = + decode_response("error id=256 msg=command\\snot\\sfound\n\r".to_string()).unwrap_err(); + + assert!(matches!( + error, + QueryError::Status { id: 256, message } if message == "command not found" + )); + } + + #[test] + fn escapes_query_values() { + assert_eq!(escape("a b|c/d\\e"), "a\\sb\\pc\\/d\\\\e"); + } + + #[test] + fn decodes_permissionlist_records() { + let response = decode_response( + "permid=1 permname=b_serverinstance_help_view permdesc=Retrieve\\sinformation\\sabout\\sServerQuery\\scommands|permid=32769 permname=i_needed_modify_power_serverinstance_help_view\r\nerror id=0 msg=ok\r\n" + .to_string(), + ) + .unwrap(); + + let permissions = records_to_permissions(&response.records).unwrap(); + assert_eq!(permissions.len(), 2); + assert_eq!(permissions[0].id, PermissionId(1)); + assert_eq!(permissions[0].name, "b_serverinstance_help_view"); + assert_eq!( + permissions[0].description, + "Retrieve information about ServerQuery commands" + ); + assert_eq!(permissions[1].id, PermissionId(32769)); + assert_eq!(permissions[1].description, ""); + } + + #[test] + fn decodes_common_serverquery_records() { + let response = decode_response( + "cid=1 pid=0 channel_order=0 channel_name=Lobby total_clients=2 channel_needed_subscribe_power=0|cid=2 pid=1 channel_order=1 channel_name=Voice\\sRoom total_clients=0 channel_needed_subscribe_power=25\r\nerror id=0 msg=ok\r\n" + .to_string(), + ) + .unwrap(); + let channels = records_to_channels(&response.records).unwrap(); + assert_eq!(channels.len(), 2); + assert_eq!(channels[0].id, ChannelId(1)); + assert_eq!(channels[1].name, "Voice Room"); + assert_eq!(channels[1].needed_subscribe_power, 25); + + let response = decode_response( + "clid=8 cid=1 client_database_id=1 client_nickname=serveradmin client_type=1 client_unique_identifier=serveradmin|clid=9 cid=2 client_database_id=42 client_nickname=Normal\\sUser client_type=0 client_unique_identifier=abc\r\nerror id=0 msg=ok\r\n" + .to_string(), + ) + .unwrap(); + let clients = records_to_clients(&response.records).unwrap(); + assert_eq!(clients.len(), 2); + assert_eq!(clients[0].id, ClientId(8)); + assert_eq!(clients[0].client_type, ClientType::Query { admin: false }); + assert_eq!(clients[1].nickname, "Normal User"); + assert_eq!(clients[1].client_type, ClientType::Normal); + + let response = decode_response( + "virtualserver_name=Test\\sServer virtualserver_platform=Linux virtualserver_version=3.13.7 virtualserver_maxclients=32 virtualserver_clientsonline=4 virtualserver_channelsonline=12 virtualserver_uptime=3600\r\nerror id=0 msg=ok\r\n" + .to_string(), + ) + .unwrap(); + let server = record_to_server_info(&response.records[0]).unwrap(); + assert_eq!(server.name, "Test Server"); + assert_eq!(server.max_clients, 32); + assert_eq!(server.clients_online, 4); + assert_eq!(server.channels_online, 12); + assert_eq!(server.uptime, 3600); + } + + #[tokio::test] + async fn executes_commands_against_mock_server() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + stream + .write_all(b"TS3\r\nWelcome to the TeamSpeak 3 ServerQuery interface\r\n") + .await + .unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "login serveradmin secret\\spass"); + stream.write_all(b"error id=0 msg=ok\r\n").await.unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "whoami"); + stream + .write_all( + b"clid=4 client_database_id=10 client_nickname=serveradmin\r\nerror id=0 msg=ok\r\n", + ) + .await + .unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "permissionlist"); + stream + .write_all( + b"permid=24 permname=b_virtualserver_select permdesc=Select\\sa\\svirtual\\sserver|permid=248 permname=i_ft_quota_mb_upload_per_client permdesc=Upload\\squota\\sper\\sclient\\sin\\sMByte\r\nerror id=0 msg=ok\r\n", + ) + .await + .unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "channellist"); + stream + .write_all( + b"cid=1 pid=0 channel_order=0 channel_name=Lobby total_clients=1 channel_needed_subscribe_power=0\r\nerror id=0 msg=ok\r\n", + ) + .await + .unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "clientlist"); + stream + .write_all( + b"clid=9 cid=1 client_database_id=42 client_nickname=Normal\\sUser client_type=0 client_unique_identifier=abc\r\nerror id=0 msg=ok\r\n", + ) + .await + .unwrap(); + + let command = read_command(&mut stream).await; + assert_eq!(command, "serverinfo"); + stream + .write_all( + b"virtualserver_name=Mock\\sServer virtualserver_platform=Linux virtualserver_version=3.13.7 virtualserver_maxclients=32 virtualserver_clientsonline=1 virtualserver_channelsonline=1 virtualserver_uptime=99\r\nerror id=0 msg=ok\r\n", + ) + .await + .unwrap(); + }); + + let mut client = QueryClient::connect(addr).await.unwrap(); + assert!(client.greeting().contains("TS3")); + client.login("serveradmin", "secret pass").await.unwrap(); + + let whoami = client.whoami().await.unwrap().unwrap(); + assert_eq!(whoami.get("clid"), Some("4")); + assert_eq!(whoami.get("client_database_id"), Some("10")); + + let permissions = client.permission_list().await.unwrap(); + assert_eq!(permissions.len(), 2); + assert_eq!(permissions[0].id, PermissionId(24)); + assert_eq!(permissions[0].name, "b_virtualserver_select"); + assert_eq!(permissions[1].id, PermissionId(248)); + + let channels = client.channel_list().await.unwrap(); + assert_eq!(channels[0].name, "Lobby"); + assert_eq!(channels[0].total_clients, 1); + + let clients = client.client_list().await.unwrap(); + assert_eq!(clients[0].nickname, "Normal User"); + assert_eq!(clients[0].database_id, ClientDbId(42)); + + let server_info = client.server_info().await.unwrap().unwrap(); + assert_eq!(server_info.name, "Mock Server"); + assert_eq!(server_info.uptime, 99); + + server.await.unwrap(); + } + + async fn read_command(stream: &mut TcpStream) -> String { + let mut data = Vec::new(); + let mut buffer = [0u8; 64]; + loop { + let len = stream.read(&mut buffer).await.unwrap(); + assert_ne!(len, 0); + data.extend_from_slice(&buffer[..len]); + let content = String::from_utf8_lossy(&data); + if content.ends_with("\n\r") || content.ends_with("\r\n") { + return content.trim_end_matches(['\n', '\r']).to_string(); + } + } + } +} diff --git a/src/tsdb/src/bookmark.rs b/src/tsdb/src/bookmark.rs index 0d9fcc3..8c05061 100644 --- a/src/tsdb/src/bookmark.rs +++ b/src/tsdb/src/bookmark.rs @@ -1,9 +1,9 @@ //! 书签管理 -use rusqlite::params; use chrono::Utc; +use rusqlite::params; -use super::{DatabaseManager, DatabaseResult, DatabaseError}; +use super::{DatabaseError, DatabaseManager, DatabaseResult}; /// 书签信息 #[derive(Debug, Clone)] @@ -64,23 +64,25 @@ impl DatabaseManager { "SELECT id, name, address, port, nickname, server_password, channel, channel_password, default_token, auto_connect, last_connected, created_at, updated_at FROM bookmarks WHERE id = ?1" )?; - let bookmark = stmt.query_row(params![id], |row| { - Ok(Bookmark { - id: row.get(0)?, - name: row.get(1)?, - address: row.get(2)?, - port: row.get(3)?, - nickname: row.get(4)?, - server_password: row.get(5)?, - channel: row.get(6)?, - channel_password: row.get(7)?, - default_token: row.get(8)?, - auto_connect: row.get::<_, i32>(9)? != 0, - last_connected: row.get(10)?, - created_at: row.get(11)?, - updated_at: row.get(12)?, + let bookmark = stmt + .query_row(params![id], |row| { + Ok(Bookmark { + id: row.get(0)?, + name: row.get(1)?, + address: row.get(2)?, + port: row.get(3)?, + nickname: row.get(4)?, + server_password: row.get(5)?, + channel: row.get(6)?, + channel_password: row.get(7)?, + default_token: row.get(8)?, + auto_connect: row.get::<_, i32>(9)? != 0, + last_connected: row.get(10)?, + created_at: row.get(11)?, + updated_at: row.get(12)?, + }) }) - }).map_err(|_| DatabaseError::NotFound(format!("书签 {} 未找到", id)))?; + .map_err(|_| DatabaseError::NotFound(format!("书签 {} 未找到", id)))?; Ok(bookmark) } @@ -92,23 +94,25 @@ impl DatabaseManager { "SELECT id, name, address, port, nickname, server_password, channel, channel_password, default_token, auto_connect, last_connected, created_at, updated_at FROM bookmarks ORDER BY name" )?; - let bookmarks = stmt.query_map([], |row| { - Ok(Bookmark { - id: row.get(0)?, - name: row.get(1)?, - address: row.get(2)?, - port: row.get(3)?, - nickname: row.get(4)?, - server_password: row.get(5)?, - channel: row.get(6)?, - channel_password: row.get(7)?, - default_token: row.get(8)?, - auto_connect: row.get::<_, i32>(9)? != 0, - last_connected: row.get(10)?, - created_at: row.get(11)?, - updated_at: row.get(12)?, - }) - })?.collect::, _>>()?; + let bookmarks = stmt + .query_map([], |row| { + Ok(Bookmark { + id: row.get(0)?, + name: row.get(1)?, + address: row.get(2)?, + port: row.get(3)?, + nickname: row.get(4)?, + server_password: row.get(5)?, + channel: row.get(6)?, + channel_password: row.get(7)?, + default_token: row.get(8)?, + auto_connect: row.get::<_, i32>(9)? != 0, + last_connected: row.get(10)?, + created_at: row.get(11)?, + updated_at: row.get(12)?, + }) + })? + .collect::, _>>()?; Ok(bookmarks) } @@ -157,10 +161,8 @@ impl DatabaseManager { /// 删除书签 pub fn delete_bookmark(&self, id: &str) -> DatabaseResult<()> { - self.connection().execute( - "DELETE FROM bookmarks WHERE id = ?1", - params![id], - )?; + self.connection() + .execute("DELETE FROM bookmarks WHERE id = ?1", params![id])?; Ok(()) } diff --git a/src/tsdb/src/config.rs b/src/tsdb/src/config.rs index 7d1e16a..21ffed9 100644 --- a/src/tsdb/src/config.rs +++ b/src/tsdb/src/config.rs @@ -1,8 +1,8 @@ //! 配置管理 +use chrono::Utc; use rusqlite::params; use rusqlite::OptionalExtension; -use chrono::Utc; use super::{DatabaseManager, DatabaseResult}; @@ -11,9 +11,9 @@ impl DatabaseManager { let conn = self.connection(); let mut stmt = conn.prepare("SELECT value FROM settings WHERE key = ?1")?; - let result = stmt.query_row(params![key], |row| { - Ok(row.get::<_, String>(0)?) - }).optional()?; + let result = stmt + .query_row(params![key], |row| Ok(row.get::<_, String>(0)?)) + .optional()?; Ok(result) } @@ -30,10 +30,8 @@ impl DatabaseManager { } pub fn delete_setting(&self, key: &str) -> DatabaseResult<()> { - self.connection().execute( - "DELETE FROM settings WHERE key = ?1", - params![key], - )?; + self.connection() + .execute("DELETE FROM settings WHERE key = ?1", params![key])?; Ok(()) } @@ -41,9 +39,9 @@ impl DatabaseManager { let conn = self.connection(); let mut stmt = conn.prepare("SELECT key, value FROM settings ORDER BY key")?; - let settings = stmt.query_map([], |row| { - Ok((row.get(0)?, row.get(1)?)) - })?.collect::, _>>()?; + let settings = stmt + .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))? + .collect::, _>>()?; Ok(settings) } diff --git a/src/tsdb/src/identity.rs b/src/tsdb/src/identity.rs index 02e8ac0..8a903e6 100644 --- a/src/tsdb/src/identity.rs +++ b/src/tsdb/src/identity.rs @@ -1,9 +1,9 @@ //! 身份管理 -use rusqlite::params; use chrono::Utc; +use rusqlite::params; -use super::{DatabaseManager, DatabaseResult, DatabaseError}; +use super::{DatabaseError, DatabaseManager, DatabaseResult}; /// 身份信息 #[derive(Debug, Clone)] @@ -46,17 +46,19 @@ impl DatabaseManager { "SELECT id, name, private_key, counter, max_counter, created_at, updated_at FROM identities WHERE id = ?1" )?; - let identity = stmt.query_row(params![id], |row| { - Ok(Identity { - id: row.get(0)?, - name: row.get(1)?, - private_key: row.get(2)?, - counter: row.get(3)?, - max_counter: row.get(4)?, - created_at: row.get(5)?, - updated_at: row.get(6)?, + let identity = stmt + .query_row(params![id], |row| { + Ok(Identity { + id: row.get(0)?, + name: row.get(1)?, + private_key: row.get(2)?, + counter: row.get(3)?, + max_counter: row.get(4)?, + created_at: row.get(5)?, + updated_at: row.get(6)?, + }) }) - }).map_err(|_| DatabaseError::NotFound(format!("身份 {} 未找到", id)))?; + .map_err(|_| DatabaseError::NotFound(format!("身份 {} 未找到", id)))?; Ok(identity) } @@ -68,23 +70,30 @@ impl DatabaseManager { "SELECT id, name, private_key, counter, max_counter, created_at, updated_at FROM identities ORDER BY name" )?; - let identities = stmt.query_map([], |row| { - Ok(Identity { - id: row.get(0)?, - name: row.get(1)?, - private_key: row.get(2)?, - counter: row.get(3)?, - max_counter: row.get(4)?, - created_at: row.get(5)?, - updated_at: row.get(6)?, - }) - })?.collect::, _>>()?; + let identities = stmt + .query_map([], |row| { + Ok(Identity { + id: row.get(0)?, + name: row.get(1)?, + private_key: row.get(2)?, + counter: row.get(3)?, + max_counter: row.get(4)?, + created_at: row.get(5)?, + updated_at: row.get(6)?, + }) + })? + .collect::, _>>()?; Ok(identities) } /// 更新身份 - pub fn update_identity(&self, id: &str, name: Option<&str>, counter: Option) -> DatabaseResult<()> { + pub fn update_identity( + &self, + id: &str, + name: Option<&str>, + counter: Option, + ) -> DatabaseResult<()> { let now = Utc::now().to_rfc3339(); if let Some(name) = name { @@ -106,10 +115,8 @@ impl DatabaseManager { /// 删除身份 pub fn delete_identity(&self, id: &str) -> DatabaseResult<()> { - self.connection().execute( - "DELETE FROM identities WHERE id = ?1", - params![id], - )?; + self.connection() + .execute("DELETE FROM identities WHERE id = ?1", params![id])?; Ok(()) } } diff --git a/src/tsdb/src/lib.rs b/src/tsdb/src/lib.rs index 93f95ca..a7a3d69 100644 --- a/src/tsdb/src/lib.rs +++ b/src/tsdb/src/lib.rs @@ -1,12 +1,12 @@ //! 数据存储 -pub mod identity; pub mod bookmark; -pub mod message; pub mod config; +pub mod identity; +pub mod message; -pub use identity::*; pub use bookmark::*; +pub use identity::*; pub use message::*; use thiserror::Error; @@ -93,7 +93,7 @@ impl DatabaseManager { value TEXT NOT NULL, updated_at TEXT NOT NULL ); - " + ", )?; Ok(()) } diff --git a/src/tsdb/src/message.rs b/src/tsdb/src/message.rs index 5912c09..26e9918 100644 --- a/src/tsdb/src/message.rs +++ b/src/tsdb/src/message.rs @@ -1,9 +1,9 @@ //! 消息管理 -use rusqlite::params; use chrono::Utc; +use rusqlite::params; -use super::{DatabaseManager, DatabaseResult, DatabaseError}; +use super::{DatabaseError, DatabaseManager, DatabaseResult}; /// 消息信息 #[derive(Debug, Clone)] @@ -62,20 +62,22 @@ impl DatabaseManager { "SELECT id, server_address, invoker_id, invoker_name, invoker_uid, target_type, target_id, message, is_read, timestamp FROM messages WHERE id = ?1" )?; - let message = stmt.query_row(params![id], |row| { - Ok(Message { - id: row.get(0)?, - server_address: row.get(1)?, - invoker_id: row.get(2)?, - invoker_name: row.get(3)?, - invoker_uid: row.get(4)?, - target_type: row.get(5)?, - target_id: row.get(6)?, - message: row.get(7)?, - is_read: row.get::<_, i32>(8)? != 0, - timestamp: row.get(9)?, + let message = stmt + .query_row(params![id], |row| { + Ok(Message { + id: row.get(0)?, + server_address: row.get(1)?, + invoker_id: row.get(2)?, + invoker_name: row.get(3)?, + invoker_uid: row.get(4)?, + target_type: row.get(5)?, + target_id: row.get(6)?, + message: row.get(7)?, + is_read: row.get::<_, i32>(8)? != 0, + timestamp: row.get(9)?, + }) }) - }).map_err(|_| DatabaseError::NotFound(format!("消息 {} 未找到", id)))?; + .map_err(|_| DatabaseError::NotFound(format!("消息 {} 未找到", id)))?; Ok(message) } @@ -92,39 +94,37 @@ impl DatabaseManager { "SELECT id, server_address, invoker_id, invoker_name, invoker_uid, target_type, target_id, message, is_read, timestamp FROM messages WHERE server_address = ?1 ORDER BY timestamp DESC LIMIT ?2 OFFSET ?3" )?; - let messages = stmt.query_map(params![server_address, limit, offset], |row| { - Ok(Message { - id: row.get(0)?, - server_address: row.get(1)?, - invoker_id: row.get(2)?, - invoker_name: row.get(3)?, - invoker_uid: row.get(4)?, - target_type: row.get(5)?, - target_id: row.get(6)?, - message: row.get(7)?, - is_read: row.get::<_, i32>(8)? != 0, - timestamp: row.get(9)?, - }) - })?.collect::, _>>()?; + let messages = stmt + .query_map(params![server_address, limit, offset], |row| { + Ok(Message { + id: row.get(0)?, + server_address: row.get(1)?, + invoker_id: row.get(2)?, + invoker_name: row.get(3)?, + invoker_uid: row.get(4)?, + target_type: row.get(5)?, + target_id: row.get(6)?, + message: row.get(7)?, + is_read: row.get::<_, i32>(8)? != 0, + timestamp: row.get(9)?, + }) + })? + .collect::, _>>()?; Ok(messages) } /// 标记消息为已读 pub fn mark_message_read(&self, id: i64) -> DatabaseResult<()> { - self.connection().execute( - "UPDATE messages SET is_read = 1 WHERE id = ?1", - params![id], - )?; + self.connection() + .execute("UPDATE messages SET is_read = 1 WHERE id = ?1", params![id])?; Ok(()) } /// 删除消息 pub fn delete_message(&self, id: i64) -> DatabaseResult<()> { - self.connection().execute( - "DELETE FROM messages WHERE id = ?1", - params![id], - )?; + self.connection() + .execute("DELETE FROM messages WHERE id = ?1", params![id])?; Ok(()) }