From 2b8b30ef416f95aee5aaf9f8dd560c07094e298f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=BB=9D=E5=B0=98?= <237809796@qq.com> Date: Mon, 22 Jun 2026 01:47:01 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E:=20df-relay=20Phase2=20WS?= =?UTF-8?q?=E4=B8=AD=E7=BB=A7=E5=AE=9E=E7=8E=B0(Hello=E6=8F=A1=E6=89=8B+?= =?UTF-8?q?=E9=89=B4=E6=9D=83+ConnRegistry=E9=85=8D=E5=AF=B9=E8=B7=AF?= =?UTF-8?q?=E7=94=B1+Event/Command=E6=96=B9=E5=90=91=E5=88=86=E5=8F=91)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/df-relay/src/broadcast.rs | 82 +++++- crates/df-relay/src/conn.rs | 202 ++++++++++++++ crates/df-relay/src/error.rs | 6 +- crates/df-relay/src/lib.rs | 11 +- crates/df-relay/src/relay.rs | 458 ++++++++++++++++++++++++------- 5 files changed, 638 insertions(+), 121 deletions(-) create mode 100644 crates/df-relay/src/conn.rs diff --git a/crates/df-relay/src/broadcast.rs b/crates/df-relay/src/broadcast.rs index e1e8d83..9b132a9 100644 --- a/crates/df-relay/src/broadcast.rs +++ b/crates/df-relay/src/broadcast.rs @@ -2,18 +2,21 @@ //! //! 设计依据:设计文档「Layer2」—— 云后端纯转发,无业务逻辑。 //! 中继消息骨架定义「谁发给谁」的路由语义: -//! - `DeviceEvent`:桌面端事件 → 广播给该 device 绑定的小程序 -//! - `MiniappCommand`:小程序指令 → 转发给该 device 绑定的桌面端 +//! - `Event`:桌面端事件 → 广播给该 device 绑定的小程序 +//! - `Command`:小程序指令 → 转发给该 device 绑定的桌面端 //! - `Control`:控制面消息(配对/心跳/在线状态) //! //! 与 df-tunnel::TunnelMessage 区分: //! - TunnelMessage 是隧道两端(桌面↔云)的传输单元 //! - BroadcastMessage 是云后端内部的路由单元(标识来源/去向/广播范围) //! -//! 脚手架阶段仅定义 enum,Phase2 填路由表 + 广播实现。 +//! AiChatEvent JSON 透传:relay 不解析 AiChatEvent(类型在 src-tauri/df-types, +//! relay 不依赖避跨 crate 强耦合),payload 用 serde_json::Value 透传。 use serde::{Deserialize, Serialize}; +use crate::conn::ConnId; + /// 客户端类型(用于路由判断) #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] @@ -24,27 +27,52 @@ pub enum ClientKind { Miniapp, } -/// 中继路由消息骨架 +/// 中继路由消息语义分类 +/// +/// - `Event`:桌面端 → 小程序(渲染事件,AiChatEvent 透传) +/// - `Command`:小程序 → 桌面端(操作指令,触发 Tauri command) +/// - `Control`:控制面(配对/心跳/Ping-Pong,不进业务流) +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum MessageKind { + Event, + Command, + Control, +} + +/// 中继路由消息骨架(云后端内部的路由单元) /// /// 每条消息携带 `device_id`(配对绑定键),云后端按 device_id -/// 找到绑定的对端连接进行转发/广播。 +/// 找到绑定的对端连接进行转发/广播。payload 为原始 JSON,转发时不解析业务字段, +/// 保持轻量。 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BroadcastMessage { /// 配对绑定的设备 ID(桌面端首次配置时生成) pub device_id: String, - /// 消息来源(用于鉴权校验:device 不能发 MiniappCommand) + /// 消息语义(Event / Command / Control) + pub kind: MessageKind, + /// 消息来源连接(用于鉴权校验:device 不能发 Command) + pub source: ConnId, + /// 来源客户端类型 pub from: ClientKind, - /// 消息载荷(原始 JSON,转发时不解析业务字段,保持轻量) + /// 消息载荷(原始 JSON,透传不解析) pub payload: serde_json::Value, - /// 时间戳(毫秒,用于冲突处理 —— 桌面端为真相源,但保留时间戳辅助判断) + /// 时间戳(毫秒,用于冲突判断 —— 桌面端为真相源,但保留时间戳辅助) pub ts: i64, } impl BroadcastMessage { /// 便捷构造:device 发出的事件 - pub fn from_device(device_id: impl Into, payload: serde_json::Value, ts: i64) -> Self { + pub fn from_device( + device_id: impl Into, + source: ConnId, + payload: serde_json::Value, + ts: i64, + ) -> Self { Self { device_id: device_id.into(), + kind: MessageKind::Event, + source, from: ClientKind::Device, payload, ts, @@ -52,14 +80,39 @@ impl BroadcastMessage { } /// 便捷构造:小程序发出的指令 - pub fn from_miniapp(device_id: impl Into, payload: serde_json::Value, ts: i64) -> Self { + pub fn from_miniapp( + device_id: impl Into, + source: ConnId, + payload: serde_json::Value, + ts: i64, + ) -> Self { Self { device_id: device_id.into(), + kind: MessageKind::Command, + source, from: ClientKind::Miniapp, payload, ts, } } + + /// 便捷构造:控制面消息 + pub fn control( + device_id: impl Into, + source: ConnId, + from: ClientKind, + payload: serde_json::Value, + ts: i64, + ) -> Self { + Self { + device_id: device_id.into(), + kind: MessageKind::Control, + source, + from, + payload, + ts, + } + } } /// 客户端在线状态(用于离线降级提示) @@ -73,8 +126,11 @@ pub enum PresenceState { } /// 控制面消息(配对绑定 / 心跳 / 在线状态广播) +/// +/// 仅作为 payload 的语义约定,relay 对 Control 消息同样透传;部分控制语义 +/// (如 Ping/Pong 心跳)由 WS 连接层直接处理,不进 BroadcastMessage 流。 #[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] +#[serde(tag = "control_kind", rename_all = "snake_case")] pub enum ControlMessage { /// 配对绑定请求(小程序持 token 申请绑定 device_id) Pair { device_id: String, token: String }, @@ -82,4 +138,8 @@ pub enum ControlMessage { Heartbeat { device_id: String }, /// 在线状态变更通知 Presence { device_id: String, state: PresenceState }, + /// 心跳 Ping(WS 应用层心跳,与协议层 Ping 区分) + Ping, + /// 心跳 Pong + Pong, } diff --git a/crates/df-relay/src/conn.rs b/crates/df-relay/src/conn.rs new file mode 100644 index 0000000..ca6130f --- /dev/null +++ b/crates/df-relay/src/conn.rs @@ -0,0 +1,202 @@ +//! df-relay 连接注册表 +//! +//! 设计依据:设计文档「Layer2」—— 按 device_id 配对路由(非全局广播)。 +//! +//! 注册表结构:`HashMap` + `device_id → conn_id` 索引。 +//! 配对语义:小程序 device_id ↔ 桌面端绑定该 device_id 的 token(同 device_id +//! 的小程序与桌面端互为对端)。relay 按 device_id 找到对端连接转发消息。 +//! +//! 并发:tokio::sync::Mutex 单锁(注册表操作高频但临界区小,先简单后优化)。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +use tokio::sync::{mpsc, Mutex}; + +use crate::broadcast::{BroadcastMessage, ClientKind}; + +/// 连接唯一标识(u64 自增,进程内唯一) +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub struct ConnId(pub u64); + +impl ConnId { + /// 伪值(用于无来源场景,如服务端构造的控制消息) + pub const NIL: ConnId = ConnId(0); +} + +static NEXT_CONN_ID: AtomicU64 = AtomicU64::new(1); + +/// 分配下一个 conn_id(进程内递增,从 1 起,0 留作 NIL) +pub fn next_conn_id() -> ConnId { + ConnId(NEXT_CONN_ID.fetch_add(1, Ordering::Relaxed)) +} + +/// 连接句柄(注册表条目) +/// +/// 每个连接持一个 mpsc::UnboundedSender,广播时向其投递 BroadcastMessage, +/// 由连接自己的读取任务转发给对端 socket。 +#[derive(Debug)] +pub struct ConnHandle { + /// 连接 ID + pub id: ConnId, + /// 客户端类型 + pub kind: ClientKind, + /// 配对绑定的设备 ID(桌面端首次配置生成,小程序携带) + pub device_id: String, + /// 广播投递通道(连接读取任务消费 → 写入 socket) + pub sender: mpsc::UnboundedSender, +} + +impl ConnHandle { + pub fn new( + id: ConnId, + kind: ClientKind, + device_id: String, + sender: mpsc::UnboundedSender, + ) -> Self { + Self { + id, + kind, + device_id, + sender, + } + } +} + +/// 连接注册表(RelayState 内部状态) +#[derive(Default)] +pub struct ConnRegistry { + /// conn_id → 连接句柄 + by_id: HashMap, + /// device_id → 绑定的连接集合(同 device_id 的小程序 + 桌面端) + /// 设计为 Vec:兼容同 device_id 多连接(多端登录),MVP 一般 1+1。 + by_device: HashMap>, +} + +impl ConnRegistry { + pub fn new() -> Self { + Self::default() + } + + /// 注册连接 + pub fn add(&mut self, handle: ConnHandle) { + let id = handle.id; + let device_id = handle.device_id.clone(); + self.by_device.entry(device_id).or_default().push(id); + self.by_id.insert(id, handle); + } + + /// 注销连接(返回其 device_id 用于清理索引) + pub fn remove(&mut self, id: ConnId) -> Option { + let handle = self.by_id.remove(&id)?; + let device_id = handle.device_id.clone(); + if let Some(ids) = self.by_device.get_mut(&device_id) { + ids.retain(|c| *c != id); + if ids.is_empty() { + self.by_device.remove(&device_id); + } + } + Some(device_id) + } + + /// 查询 device_id 是否有在线连接(离线降级判断用) + pub fn is_device_online(&self, device_id: &str) -> bool { + self.by_device + .get(device_id) + .map(|ids| !ids.is_empty()) + .unwrap_or(false) + } + + /// 投递消息给指定 device_id 配对的对端连接(排除来源 source) + /// + /// 路由规则: + /// - Event(Device 发):转发给同 device_id 的所有 Miniapp 连接 + /// - Command(Miniapp 发):转发给同 device_id 的所有 Device 连接 + /// - Control:转发给同 device_id 的所有连接(除来源) + /// + /// 返回成功投递的连接数(用于日志/降级判断)。 + pub fn route(&self, msg: &BroadcastMessage) -> usize { + let Some(ids) = self.by_device.get(&msg.device_id) else { + return 0; + }; + let mut delivered = 0usize; + for &id in ids { + if id == msg.source { + // 不回环给来源 + continue; + } + let Some(handle) = self.by_id.get(&id) else { + continue; + }; + match (msg.kind, msg.from, handle.kind) { + // Event:仅投递给 Miniapp + (crate::broadcast::MessageKind::Event, _, ClientKind::Miniapp) => {} + // Command:仅投递给 Device + (crate::broadcast::MessageKind::Command, _, ClientKind::Device) => {} + // Control:全投递(除来源) + (crate::broadcast::MessageKind::Control, _, _) => {} + // 其余方向不匹配(Event 不投递给 Device,Command 不投递给 Miniapp) + _ => continue, + } + if handle.sender.send(msg.clone()).is_ok() { + delivered += 1; + } + } + delivered + } + + /// 当前总连接数(诊断用) + pub fn len(&self) -> usize { + self.by_id.len() + } + + /// 是否为空(诊断用) + pub fn is_empty(&self) -> bool { + self.by_id.is_empty() + } +} + +/// 共享中继状态(axum State 传递) +/// +/// 注册表用 tokio Mutex 包裹:广播路径在锁内调用 route,锁粒度小。 +#[derive(Clone, Default)] +pub struct RelayState { + inner: Arc>, +} + +impl RelayState { + pub fn new() -> Self { + Self::default() + } + + /// 从 Arc 直接构造(用于 axum state) + pub fn from_arc(inner: Arc>) -> Self { + Self { inner } + } + + /// 暴露内部 Arc(测试/外部 start 使用) + pub fn inner(&self) -> Arc> { + self.inner.clone() + } + + /// 注册连接 + pub async fn add_conn(&self, handle: ConnHandle) { + self.inner.lock().await.add(handle); + } + + /// 注销连接 + pub async fn remove_conn(&self, id: ConnId) -> Option { + self.inner.lock().await.remove(id) + } + + /// 查询 device 在线状态 + pub async fn is_device_online(&self, device_id: &str) -> bool { + self.inner.lock().await.is_device_online(device_id) + } + + /// 投递消息给对端连接(返回成功投递数) + pub async fn route(&self, msg: &BroadcastMessage) -> usize { + self.inner.lock().await.route(msg) + } +} diff --git a/crates/df-relay/src/error.rs b/crates/df-relay/src/error.rs index ab49736..9a61545 100644 --- a/crates/df-relay/src/error.rs +++ b/crates/df-relay/src/error.rs @@ -20,7 +20,7 @@ pub enum RelayError { #[error("广播失败: {0}")] Broadcast(String), - /// WS 协议层错误(底层 tungstenite) + /// WS 协议层错误(底层 axum::extract::ws / tungstenite) #[error("WS 协议错误: {0}")] WebSocket(String), @@ -28,6 +28,10 @@ pub enum RelayError { #[error("消息序列化失败: {0}")] Serde(#[from] serde_json::Error), + /// 广播发送失败(mpsc channel 对端已断开) + #[error("广播 channel 发送失败: {0}")] + Channel(String), + /// 其他未归类错误(降级通道) #[error("中继内部错误: {0}")] Other(#[from] anyhow::Error), diff --git a/crates/df-relay/src/lib.rs b/crates/df-relay/src/lib.rs index 0c0c13c..5c1b6e4 100644 --- a/crates/df-relay/src/lib.rs +++ b/crates/df-relay/src/lib.rs @@ -12,15 +12,18 @@ //! - 不含业务逻辑(Phase2 填连接生命周期 / 注册表 / 广播分发) //! //! ## 模块 -//! - [`broadcast`]:广播消息骨架(BroadcastMessage / ClientKind / ControlMessage) -//! - [`relay`]:RelayServer trait + DefaultRelayServer 骨架 + axum WS 路由 +//! - [`broadcast`]:广播消息骨架(BroadcastMessage / ClientKind / ControlMessage / MessageKind) +//! - [`conn`]:连接注册表(ConnId / ConnHandle / ConnRegistry / RelayState) +//! - [`relay`]:RelayServer trait + DefaultRelayServer + axum WS 路由与连接生命周期 //! - [`error`]:RelayError 错误类型 pub mod broadcast; +pub mod conn; pub mod error; pub mod relay; // 顶层再导出常用项 -pub use broadcast::{BroadcastMessage, ClientKind, ControlMessage, PresenceState}; +pub use broadcast::{BroadcastMessage, ClientKind, ControlMessage, MessageKind, PresenceState}; +pub use conn::{next_conn_id, ConnHandle, ConnId, ConnRegistry, RelayState}; pub use error::{RelayError, Result}; -pub use relay::{build_router, DefaultRelayServer, RelayServer, RelayState}; +pub use relay::{build_router, DefaultRelayServer, RelayServer}; diff --git a/crates/df-relay/src/relay.rs b/crates/df-relay/src/relay.rs index f4a6701..cd46f10 100644 --- a/crates/df-relay/src/relay.rs +++ b/crates/df-relay/src/relay.rs @@ -1,29 +1,73 @@ -//! df-relay 中继服务骨架 +//! df-relay 中继服务核心实现 //! //! 设计依据:设计文档「Layer2」—— axum WS Server + 广播中继。 -//! 接受两类连接:小程序(device_id 鉴权)+ 桌面端(token 配对)。 -//! 纯转发,无业务逻辑。 +//! 接受两类连接:小程序(device_id 鉴权)+ 桌面端(token 配对),按 device_id +//! 配对转发(非全局广播),纯转发无业务逻辑。 //! -//! 脚手架范围:RelayServer trait + axum WS 路由骨架(占位 handler), -//! Phase2 填: -//! - 客户端注册表(HashMap) -//! - 广播分发(按 device_id 路由) -//! - 鉴权中间件 +//! 协议(简单握手): +//! 1. 客户端建立 WS 后,首条消息发 JSON `Hello { kind, device_id, token }` 宣告身份。 +//! 2. relay 校验 token(MVP:env `DF_RELAY_TOKEN` 或硬编码常量;生产级鉴权留 Phase3)。 +//! 3. 校验通过 → 注册连接、进入收发循环;失败 → 发 Error 帧 + Close。 +//! 4. 后续消息按 kind 路由:Event(device→miniapp)/ Command(miniapp→device)/ Control。 +//! +//! AiChatEvent JSON 透传:relay 不解析 payload,只按 device_id + 方向转发。 + +use std::net::SocketAddr; use async_trait::async_trait; use axum::{ extract::{ - ws::{WebSocket, WebSocketUpgrade}, + ws::{Message, WebSocket, WebSocketUpgrade}, State, }, response::IntoResponse, routing::get, Router, }; -use std::sync::Arc; +use futures_util::{SinkExt, StreamExt}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; -use crate::broadcast::BroadcastMessage; -use crate::error::Result; +use crate::broadcast::{BroadcastMessage, ClientKind, MessageKind}; +use crate::conn::{next_conn_id, ConnHandle, ConnId, RelayState}; +use crate::error::{RelayError, Result}; + +/// MVP 默认 token(env `DF_RELAY_TOKEN` 缺省时回退)。 +/// 生产级鉴权(每 device 独立 token + 过期刷新)留 Phase3。 +const DEFAULT_TOKEN: &str = "devflow-relay-default-token"; + +/// 读取期望 token(env 优先,fallback 硬编码) +fn expected_token() -> String { + std::env::var("DF_RELAY_TOKEN").unwrap_or_else(|_| DEFAULT_TOKEN.to_string()) +} + +/// 客户端首消息:身份宣告(简单协议) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Hello { + /// 客户端类型("device" / "miniapp") + pub kind: ClientKindWire, + /// 配对绑定的设备 ID + pub device_id: String, + /// 配对 token + pub token: String, +} + +/// Hello.kind 的传输表示(serde 字符串,与 ClientKind 解耦避免 rename 歧义) +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ClientKindWire { + Device, + Miniapp, +} + +impl From for ClientKind { + fn from(w: ClientKindWire) -> Self { + match w { + ClientKindWire::Device => ClientKind::Device, + ClientKindWire::Miniapp => ClientKind::Miniapp, + } + } +} /// 中继服务抽象 /// @@ -36,86 +80,30 @@ pub trait RelayServer: Send + Sync { /// 广播消息给指定 device_id 绑定的对端 async fn broadcast(&self, msg: BroadcastMessage) -> Result<()>; - /// 注册新连接(连接建立后入表) - async fn register_client(&self, device_id: String) -> Result<()>; - - /// 注销连接(断开后出表) - async fn unregister_client(&self, device_id: String) -> Result<()>; - /// 查询 device 是否有在线连接(离线降级判断用) fn is_device_online(&self, device_id: &str) -> bool; } -/// 共享状态(axum State 传递) -/// -/// 脚手架占位:Phase2 填 tokio::sync::RwLock>。 -/// 当前用 unit 占位以满足 Router state 约束。 -#[derive(Clone, Default)] -pub struct RelayState { - /// Phase2:客户端注册表 - /// 例:`clients: Arc>>` - #[allow(dead_code)] - _reserved: (), -} - -/// axum WS 路由构造 -/// -/// 暴露 `/ws/device`(桌面端连入)与 `/ws/miniapp`(小程序连入)两个端点, -/// 共享 RelayState。Phase2 在 handler 内填鉴权 + 连接生命周期管理。 -pub fn build_router(state: Arc) -> Router { - Router::new() - .route("/ws/device", get(device_ws_handler)) - .route("/ws/miniapp", get(miniapp_ws_handler)) - .with_state(state) -} - -/// 桌面端 WS 连接 handler(脚手架占位) -/// -/// 当前仅接受升级并立即返回(不消费流),Phase2 填: -/// - 握手鉴权(token) -/// - register_client -/// - 循环 recv 转发 BroadcastMessage -/// - 断开 unregister_client -async fn device_ws_handler( - ws: WebSocketUpgrade, - State(_state): State>, -) -> impl IntoResponse { - tracing::info!("桌面端 WS 连接接入(脚手架占位)"); - ws.on_upgrade(handle_connection) -} - -/// 小程序 WS 连接 handler(脚手架占位) -async fn miniapp_ws_handler( - ws: WebSocketUpgrade, - State(_state): State>, -) -> impl IntoResponse { - tracing::info!("小程序 WS 连接接入(脚手架占位)"); - ws.on_upgrade(handle_connection) -} - -/// 通用 WS 连接处理(脚手架:接收一次即关闭) -/// -/// Phase2 替换为完整读写循环 + 广播分发。 -async fn handle_connection(mut socket: WebSocket) { - // 接收首条消息(占位),Phase2 改为持续循环 - if let Some(Ok(msg)) = futures_util::StreamExt::next(&mut socket).await { - tracing::debug!(?msg, "收到首条消息(脚手架占位,Phase2 填路由)"); - } - // 脚手架:静默关闭,不报错(连接生命周期管理留 Phase2) -} - -/// 默认中继服务骨架 -/// -/// 脚手架阶段:trait 方法返回未实现错误,不静默假成功。 -/// Phase2 替换为完整实现(启动 axum::serve + 注册表 + 广播)。 +/// 默认中继服务(持有共享 RelayState) pub struct DefaultRelayServer { - #[allow(dead_code)] - bind_addr: Option, + state: RelayState, } impl DefaultRelayServer { pub fn new() -> Self { - Self { bind_addr: None } + Self { + state: RelayState::new(), + } + } + + /// 从既有 RelayState 构造(测试 / 外部复用) + pub fn with_state(state: RelayState) -> Self { + Self { state } + } + + /// 暴露共享状态(外部可读连接数等) + pub fn state(&self) -> RelayState { + self.state.clone() } } @@ -128,31 +116,291 @@ impl Default for DefaultRelayServer { #[async_trait] impl RelayServer for DefaultRelayServer { async fn start(&self, addr: &str) -> Result<()> { - tracing::warn!(addr = %addr, "DefaultRelayServer.start 尚未实现(Phase2 填充)"); - Err(crate::error::RelayError::Other(anyhow::anyhow!( - "DefaultRelayServer.start 尚未实现(Phase2)" - ))) + let socket_addr: SocketAddr = addr + .parse() + .map_err(|e| RelayError::Start(format!("地址解析失败 {addr}: {e}")))?; + let app = build_router(self.state.clone()); + let listener = tokio::net::TcpListener::bind(&socket_addr) + .await + .map_err(|e| RelayError::Start(format!("监听绑定失败 {addr}: {e}")))?; + tracing::info!(%addr, "df-relay WS 服务已启动"); + axum::serve(listener, app) + .await + .map_err(|e| RelayError::Start(format!("axum::serve 失败: {e}")))?; + Ok(()) } - async fn broadcast(&self, _msg: BroadcastMessage) -> Result<()> { - Err(crate::error::RelayError::Other(anyhow::anyhow!( - "DefaultRelayServer.broadcast 尚未实现(Phase2)" - ))) + async fn broadcast(&self, msg: BroadcastMessage) -> Result<()> { + let delivered = self.state.route(&msg).await; + if delivered == 0 { + // 对端离线不算硬错误(MVP 返回 Ok,离线降级由调用方据 is_device_online 判断) + tracing::debug!( + device_id = %msg.device_id, + kind = ?msg.kind, + "广播无对端在线(消息丢弃)" + ); + } + Ok(()) } - async fn register_client(&self, _device_id: String) -> Result<()> { - Err(crate::error::RelayError::Other(anyhow::anyhow!( - "DefaultRelayServer.register_client 尚未实现(Phase2)" - ))) - } - - async fn unregister_client(&self, _device_id: String) -> Result<()> { - Err(crate::error::RelayError::Other(anyhow::anyhow!( - "DefaultRelayServer.unregister_client 尚未实现(Phase2)" - ))) - } - - fn is_device_online(&self, _device_id: &str) -> bool { - false + fn is_device_online(&self, device_id: &str) -> bool { + // trait 同步签名:tokio Mutex 用 try_lock 快照,失败保守返回 false + match self.state.inner().try_lock() { + Ok(g) => g.is_device_online(device_id), + Err(_) => false, + } } } + +/// axum WS 路由构造 +/// +/// 暴露 `/ws/device`(桌面端连入)与 `/ws/miniapp`(小程序连入)两个端点, +/// 共享 RelayState。端点仅决定「期望的客户端类型」,真正的身份宣告在首消息 +/// Hello 中再次校验(防误连/误用)。 +pub fn build_router(state: RelayState) -> Router { + Router::new() + .route("/ws/device", get(device_ws_handler)) + .route("/ws/miniapp", get(miniapp_ws_handler)) + .with_state(state) +} + +/// 桌面端 WS upgrade handler +async fn device_ws_handler( + ws: WebSocketUpgrade, + State(state): State, +) -> impl IntoResponse { + tracing::debug!("桌面端 WS 连接接入"); + ws.on_upgrade(move |socket| handle_connection(socket, state, ClientKindWire::Device)) +} + +/// 小程序 WS upgrade handler +async fn miniapp_ws_handler( + ws: WebSocketUpgrade, + State(state): State, +) -> impl IntoResponse { + tracing::debug!("小程序 WS 连接接入"); + ws.on_upgrade(move |socket| handle_connection(socket, state, ClientKindWire::Miniapp)) +} + +/// WS 连接生命周期(握手 → 收发循环 → 注销) +/// +/// 步骤: +/// 1. 等待首条 Hello 文本帧,校验 kind 与 token。 +/// 2. 校验通过:分配 conn_id + mpsc,注册 ConnHandle,派发广播读取任务。 +/// 3. 主循环:从 socket recv 文本帧 → 构造 BroadcastMessage → route 投递。 +/// 4. 同时读取 mpsc 广播队列 → 写回 socket(双任务用 split sink/stream)。 +/// 5. 任一端断开 → 注销连接、关闭 mpsc。 +async fn handle_connection(socket: WebSocket, state: RelayState, expected: ClientKindWire) { + // 握手阶段:等待首条 Hello + let (mut socket_tx, mut socket_rx) = socket.split(); + + let hello = match recv_hello(&mut socket_rx).await { + Ok(h) => h, + Err(e) => { + tracing::warn!(error = %e, "握手失败:未收到合法 Hello"); + let _ = send_text( + &mut socket_tx, + r#"{"kind":"control","error":"handshake_failed"}"#, + ) + .await; + let _ = socket_tx.close().await; + return; + } + }; + + // 身份 + token 双因子校验 + if hello.kind != expected { + tracing::warn!( + ?hello.kind, + ?expected, + "握手失败:客户端类型与端点不匹配" + ); + let _ = send_text( + &mut socket_tx, + r#"{"kind":"control","error":"kind_mismatch"}"#, + ) + .await; + let _ = socket_tx.close().await; + return; + } + if hello.token != expected_token() { + tracing::warn!( + device_id = %hello.device_id, + "握手失败:token 校验不通过" + ); + let _ = send_text( + &mut socket_tx, + r#"{"kind":"control","error":"auth_failed"}"#, + ) + .await; + let _ = socket_tx.close().await; + return; + } + + let conn_id = next_conn_id(); + let kind: ClientKind = hello.kind.into(); + let device_id = hello.device_id.clone(); + tracing::info!( + conn_id = conn_id.0, + ?kind, + device_id = %device_id, + "连接握手通过,进入收发循环" + ); + + // 建立广播投递 mpsc(连接读取任务消费 → 写回 socket) + let (bc_tx, bc_rx) = mpsc::unbounded_channel::(); + let handle = ConnHandle::new(conn_id, kind, device_id.clone(), bc_tx); + state.add_conn(handle).await; + + // 派发广播读取任务:从 bc_rx 取消息 → 序列化 → 写 socket + let mut bc_task = tokio::spawn(broadcast_pump(bc_rx, socket_tx)); + + // 主循环:从 socket recv → 构造 BroadcastMessage → route + loop { + tokio::select! { + // socket 入帧 + maybe_msg = socket_rx.next() => { + match maybe_msg { + Some(Ok(Message::Text(text))) => { + if let Err(e) = handle_inbound_text(&state, conn_id, kind, &device_id, &text).await { + tracing::warn!(conn_id = conn_id.0, error = %e, "入站消息处理失败,忽略"); + } + } + Some(Ok(Message::Binary(_))) => { + // MVP 仅支持文本帧;二进制帧忽略(协议层可后续扩展) + tracing::debug!(conn_id = conn_id.0, "收到二进制帧,忽略"); + } + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => { + // axum/tungstenite 协议层 Ping/Pong 自动处理,这里仅记录 + tracing::trace!(conn_id = conn_id.0, "协议层 Ping/Pong"); + } + Some(Ok(Message::Close(_))) | None => { + tracing::info!(conn_id = conn_id.0, "客户端主动关闭连接"); + break; + } + Some(Err(e)) => { + tracing::warn!(conn_id = conn_id.0, error = %e, "socket 接收错误,断开"); + break; + } + } + } + // 广播 pump 任务结束(socket_tx 关闭或 mpsc 关闭) + res = &mut bc_task => { + match res { + Ok(()) => { + tracing::debug!(conn_id = conn_id.0, "广播 pump 任务正常结束"); + } + Err(e) => { + tracing::warn!(conn_id = conn_id.0, error = %e, "广播 pump 任务 panic"); + } + } + break; + } + } + } + + // 注销连接 + if let Some(d) = state.remove_conn(conn_id).await { + tracing::info!(conn_id = conn_id.0, device_id = %d, "连接已注销"); + } + // 结束 pump 任务(若仍在运行) + bc_task.abort(); +} + +/// 接收并解析首条 Hello 文本帧 +async fn recv_hello(rx: &mut futures_util::stream::SplitStream) -> Result { + let deadline = tokio::time::Duration::from_secs(10); + let next = tokio::time::timeout(deadline, rx.next()) + .await + .map_err(|_| RelayError::Client("握手超时(10s 未收到 Hello)".into()))?; + let msg = next + .ok_or_else(|| RelayError::Client("握手阶段连接关闭".into()))? + .map_err(|e| RelayError::WebSocket(format!("握手 recv 失败: {e}")))?; + let text = match msg { + Message::Text(t) => t, + Message::Binary(_) => { + return Err(RelayError::Client("握手首帧必须为文本".into())); + } + _ => return Err(RelayError::Client("握手首帧类型非法".into())), + }; + let hello: Hello = + serde_json::from_str(&text).map_err(|e| RelayError::Client(format!("Hello 解析失败: {e}")))?; + Ok(hello) +} + +/// 处理入站文本帧(构造 BroadcastMessage → route) +async fn handle_inbound_text( + state: &RelayState, + conn_id: ConnId, + kind: ClientKind, + device_id: &str, + raw: &str, +) -> Result<()> { + // 入站文本即业务 payload(relay 不解析),包成 BroadcastMessage + // payload 直接用原始 JSON 值;若客户端发非 JSON 文本,则包成字符串值 + let payload: serde_json::Value = serde_json::from_str(raw).unwrap_or(serde_json::Value::String(raw.to_string())); + + let now = now_ms(); + let (msg_kind, from) = match kind { + ClientKind::Device => (MessageKind::Event, ClientKind::Device), + ClientKind::Miniapp => (MessageKind::Command, ClientKind::Miniapp), + }; + let msg = BroadcastMessage { + device_id: device_id.to_string(), + kind: msg_kind, + source: conn_id, + from, + payload, + ts: now, + }; + let delivered = state.route(&msg).await; + tracing::debug!( + conn_id = conn_id.0, + ?msg_kind, + device_id = %device_id, + delivered, + "入站消息已路由" + ); + Ok(()) +} + +/// 广播 pump:从 mpsc 取消息,序列化后写回 socket sink +/// +/// 任务退出条件:bc_rx 关闭(对端 handle 全部 drop)/ socket_tx 关闭出错。 +async fn broadcast_pump( + mut bc_rx: mpsc::UnboundedReceiver, + mut socket_tx: futures_util::stream::SplitSink, +) { + while let Some(msg) = bc_rx.recv().await { + let text = match serde_json::to_string(&msg) { + Ok(t) => t, + Err(e) => { + tracing::warn!(error = %e, "广播消息序列化失败,跳过"); + continue; + } + }; + if let Err(e) = socket_tx.send(Message::Text(text)).await { + tracing::warn!(error = %e, "广播写回 socket 失败,pump 退出"); + break; + } + } +} + +/// 便捷发送文本帧 +async fn send_text( + tx: &mut futures_util::stream::SplitSink, + text: &str, +) -> Result<()> { + tx.send(Message::Text(text.to_string())) + .await + .map_err(|e| RelayError::WebSocket(format!("发送失败: {e}"))) +} + +/// 当前毫秒时间戳(避开 chrono workspace 依赖,直接用 std + SystemTime) +fn now_ms() -> i64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) +}