Files
qiming/qimingclaw/specs/agent-client/vendorsdoc/rustdesk-server.md

31 KiB
Raw Permalink Blame History

rustdesk-server

项目概述

RustDesk 官方服务器端程序,包含 hbbs信令服务器和 hbbr中继服务器

本地路径: vendors/rustdesk-server

目录结构

rustdesk-server/
├── src/
│   ├── lib.rs                    # 库入口
│   ├── rendezvous_server.rs      #  rendezvous 服务器(核心)
│   ├── relay_server.rs           # 中继服务器
│   ├── peer.rs                   # 对等节点管理
│   ├── database.rs               # SQLite数据库操作
│   ├── message.rs                # 消息处理
│   ├── heartbeat.rs              # 心跳检测
│   └── hbbr.rs                   # 中继服务器入口
├── Cargo.toml
└── libs/
    └── hbb_common/               # 通用库
        ├── src/
        │   ├── lib.rs
        │   ├── message_proto/    # protobuf 消息定义
        │   ├── net/              # 网络工具
        │   └── util/             # 工具函数
        └── Cargo.toml

核心依赖

[dependencies]
hbb_common = { path = "libs/hbb_common" }

# Web 框架
axum = "0.7"
tokio = { version = "1.0", features = ["full"] }

# 数据库
sqlx = { version = "0.7", features = ["sqlite", "runtime-tokio"] }

# WebSocket
tokio-tungstenite = "0.23"

# 命令行
clap = "4.0"

# 配置
rust-ini = "0.19"

# 认证
jsonwebtoken = "9.0"
bcrypt = "0.15"

# 工具
thiserror = "1.0"
anyhow = "1.0"
chrono = "0.4"
parking_lot = "0.12"

hbbs - 信令服务器

核心结构体

// src/rendezvous_server.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex, RwLock};
use tokio::sync::mpsc;
use futures::Sink;

pub struct RendezvousServer {
    // TCP 打洞连接池
    tcp_punch: Arc<Mutex<HashMap<SocketAddr, Sink>>>,

    // 节点映射 (ID -> PeerInfo)
    pm: Arc<PeerMap>,

    // 消息发送通道
    tx: Sender,

    // 中继服务器列表
    relay_servers: Arc<RelayServers>,

    // 内部状态
    inner: Arc<Inner>,
}

// 内部状态
struct Inner {
    // 监听地址
    addr: SocketAddr,
    // 运行状态
    running: AtomicBool,
    // 配置
    config: Config,
}

// 节点映射
pub(crate) struct PeerMap {
    // 内存缓存 (ID -> PeerInfo)
    map: Arc<RwLock<HashMap<String, LockPeer>>>,

    // SQLite 持久化
    db: database::Database,
}

// 锁住的 Peer带超时
type LockPeer = Mutex<Option<Peer>>;

pub struct Peer {
    pub id: String,              // 节点 ID
    pub uuid: String,            // UUID
    pub pk: Vec<u8>,             # 公钥
    pub options: HashMap<String, String>,  // 选项
    pub created_at: chrono::DateTime<Utc>,
    pub user: Option<String>,    // 关联用户
    pub status: i32,             // 状态
    pub note: String,
    pub info: String,
}

数据库模型 (database.rs)

// src/database.rs

use sqlx::{SqlitePool, Row};

pub struct Database {
    pool: SqlitePool,
}

impl Database {
    pub async fn new(path: &str) -> Result<Self> {
        let pool = SqlitePool::connect(path).await?;
        Ok(Self { pool })
    }

    /// 初始化数据库表
    pub async fn init(&self) -> Result<()> {
        sqlx::query!(
            r#"
            CREATE TABLE IF NOT EXISTS peer (
                guid blob PRIMARY KEY,
                id varchar(100) NOT NULL,
                uuid blob NOT NULL,
                pk blob NOT NULL,
                created_at datetime,
                user blob,
                status tinyint,
                note varchar(300),
                info text
            )
            "#
        ).execute(&self.pool).await?;

        // 创建索引
        sqlx::query!("CREATE UNIQUE INDEX IF NOT EXISTS idx_peer_id ON peer(id)").execute(&self.pool).await?;
        sqlx::query!("CREATE INDEX IF NOT EXISTS idx_peer_user ON peer(user)").execute(&self.pool).await?;
        sqlx::query!("CREATE INDEX IF NOT EXISTS idx_peer_created ON peer(created_at)").execute(&self.pool).await?;
        sqlx::query!("CREATE INDEX IF NOT EXISTS idx_peer_status ON peer(status)").execute(&self.pool).await?;

        Ok(())
    }

    /// 根据 ID 查找节点
    pub async fn get_by_id(&self, id: &str) -> Result<Option<Peer>> {
        let row = sqlx::query("SELECT * FROM peer WHERE id = ?")
            .bind(id)
            .fetch_optional(&self.pool)
            .await?;

        row.map(|r| Peer {
            id: r.try_get("id")?,
            uuid: r.try_get("uuid")?,
            pk: r.try_get("pk")?,
            // ... 其他字段
        }).transpose()
    }

    /// 保存节点
    pub async fn save_peer(&self, peer: &Peer) -> Result<()> {
        sqlx::query!(
            r#"
            INSERT OR REPLACE INTO peer (guid, id, uuid, pk, created_at, user, status, note, info)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            "#,
            uuid::Uuid::new_v4().as_bytes(),
            peer.id,
            peer.uuid,
            peer.pk,
            peer.created_at,
            peer.user,
            peer.status,
            peer.note,
            peer.info,
        ).execute(&self.pool).await?;

        Ok(())
    }

    /// 列出所有节点
    pub async fn list_peers(&self, limit: i32, offset: i32) -> Result<Vec<Peer>> {
        sqlx::query_as!(Peer, "SELECT * FROM peer ORDER BY created_at DESC LIMIT ? OFFSET ?", limit, offset)
            .fetch_all(&self.pool)
            .await
    }
}

协议消息 (rendezvous_proto)

// 注册节点
message RegisterPeer {
    string id = 1;              // 节点 ID
    bytes uuid = 2;             // UUID
    bytes pk = 3;               // 公钥
    map<string, string> options = 4;  // 选项
}

// 注册公钥
message RegisterPk {
    bytes pk = 1;               // 公钥
}

// 打洞请求
message PunchHoleRequest {
    string target_id = 1;       // 目标节点 ID
    ConnectionType conn_type = 2;  // 连接类型
}

// 打洞响应
message PunchHoleResponse {
    bool success = 1;
    SocketAddr relay_server = 2;  // 中继服务器地址
    bytes challenge = 3;         // 挑战数据
}

// 连接类型
enum ConnectionType {
    CONNECTION_TYPE_NONE = 0;
    CONNECTION_TYPE_TCP = 1;
    CONNECTION_TYPE_KCP = 2;
    CONNECTION_TYPE_WEBSOCKET = 3;
}

// 请求中继
message RequestRelay {
    string target_id = 1;
    bytes token = 2;
}

// 中继响应
message RelayResponse {
    bytes relay_server = 1;
}

消息处理流程

// src/message.rs

pub async fn handle_message(
    msg: &Message,
    peer_map: &Arc<PeerMap>,
    tx: &Sender,
) -> Result<Option<Message>> {
    match msg.type_ {
        MessageType::RegisterPeer => {
            let req = msg.payload.register_peer()?;
            handle_register_peer(peer_map, &req, tx).await
        }
        MessageType::RegisterPk => {
            let req = msg.payload.register_pk()?;
            handle_register_pk(peer_map, &req, tx).await
        }
        MessageType::PunchHoleRequest => {
            let req = msg.payload.punch_hole_request()?;
            handle_punch_hole(peer_map, &req, tx).await
        }
        MessageType::RequestRelay => {
            let req = msg.payload.request_relay()?;
            handle_request_relay(peer_map, &req, tx).await
        }
        _ => Ok(None)
    }
}

/// 处理节点注册
async fn handle_register_peer(
    peer_map: &Arc<PeerMap>,
    req: &RegisterPeer,
    tx: &Sender,
) -> Result<Option<Message>> {
    // 1. 验证 ID 格式
    if req.id.is_empty() || req.id.len() > 100 {
        return Err(Error::InvalidId);
    }

    // 2. 检查 ID 是否已被使用
    if let Some(existing) = peer_map.get_by_id(&req.id).await? {
        if existing.uuid != req.uuid {
            return Err(Error::IdAlreadyInUse);
        }
    }

    // 3. 保存到数据库
    let peer = Peer {
        id: req.id.clone(),
        uuid: req.uuid.clone(),
        pk: req.pk.clone(),
        options: req.options.clone(),
        created_at: Utc::now(),
        user: None,
        status: 1,
        note: String::new(),
        info: String::new(),
    };
    peer_map.save_peer(&peer).await?;

    // 4. 广播新节点上线
    tx.send(MessageType::PeerOnline as i32, &peer.id).await;

    Ok(Some(Message {
        type_: MessageType::RegisterPeerResponse as i32,
        payload: Some(RegisterPeerResponse { success: true }.into()),
        ..Default::default()
    }))
}

/// 处理打洞请求
async fn handle_punch_hole(
    peer_map: &Arc<PeerMap>,
    req: &PunchHoleRequest,
    tx: &Sender,
) -> Result<Option<Message>> {
    // 1. 查找目标节点
    let target = peer_map.get_by_id(&req.target_id).await?;

    match target {
        Some(peer) => {
            // 2. 获取目标连接信息
            let peer_info = peer_map.get_connection_info(&req.target_id).await?;

            // 3. 发送打洞消息给目标
            let punch_msg = Message {
                type_: MessageType::PunchHole as i32,
                payload: Some(PunchHole {
                    from_id: req.from_id.clone(),
                    conn_type: req.conn_type,
                }.into()),
                ..Default::default()
            };
            tx.send_to(&req.target_id, &punch_msg).await;

            // 4. 返回目标信息
            Ok(Some(Message {
                type_: MessageType::PunchHoleResponse as i32,
                payload: Some(PunchHoleResponse {
                    success: true,
                    relay_server: peer_info.relay_server,
                    ..Default::default()
                }.into()),
                ..Default::default()
            }))
        }
        None => {
            // 目标不在线,返回中继服务器地址
            Ok(Some(Message {
                type_: MessageType::PunchHoleResponse as i32,
                payload: Some(PunchHoleResponse {
                    success: false,
                    relay_server: get_relay_server(),
                    ..Default::default()
                }.into()),
                ..Default::default()
            }))
        }
    }
}

hbbr - 中继服务器

中继服务器架构

// src/relay_server.rs

use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use futures::StreamExt;

pub struct RelayServer {
    // 监听地址
    listen_addr: SocketAddr,

    // 带宽限制
    speed_limit: Arc<SpeedLimit>,

    // 连接管理
    connections: Arc<ConnectionManager>,

    // 配置
    config: RelayConfig,
}

pub struct RelayConfig {
    pub max_connections: usize,
    pub max_bandwidth: u64,       // 最大带宽 (bytes/s)
    pub connection_timeout: Duration,
}

impl RelayServer {
    pub async fn start(&self) -> Result<()> {
        let listener = TcpListener::bind(self.listen_addr).await?;

        tracing::info!("Relay server listening on {}", self.listen_addr);

        loop {
            let (stream, addr) = listener.accept().await?;
            let speed_limit = self.speed_limit.clone();
            let connections = self.connections.clone();

            tokio::spawn(async move {
                if let Err(e) = self.handle_connection(stream, addr, speed_limit, connections).await {
                    tracing::error!("Relay connection error: {}", e);
                }
            });
        }
    }

    async fn handle_connection(
        &self,
        stream: TcpStream,
        addr: SocketAddr,
        speed_limit: Arc<SpeedLimit>,
        connections: Arc<ConnectionManager>,
    ) -> Result<()> {
        // 1. 认证
        let (mut reader, writer) = stream.into_split();
        let auth = self.authenticate(&mut reader).await?;

        // 2. 注册连接
        let conn_id = connections.add(addr, auth).await?;

        // 3. 数据中继循环
        let mut buf = vec![0u8; 65536];
        loop {
            // 应用带宽限制
            speed_limit.wait_if_needed().await;

            // 读取数据
            let n = reader.read(&mut buf).await?;
            if n == 0 {
                break;
            }

            // 转发到目标
            connections.forward(&conn_id, &buf[..n]).await?;
        }

        // 4. 清理连接
        connections.remove(&conn_id).await;

        Ok(())
    }
}

/// 带宽限制
struct SpeedLimit {
    // 已使用的带宽
    used: AtomicU64,
    // 限制
    limit: u64,
    // 窗口开始时间
    window_start: AtomicU64,
}

impl SpeedLimit {
    pub async fn wait_if_needed(&self) {
        let now = current_time_millis();
        let window_start = self.window_start.load(Ordering::SeqCst);

        if now - window_start > 1000 {
            // 新窗口,重置计数器
            self.window_start.store(now, Ordering::SeqCst);
            self.used.store(0, Ordering::SeqCst);
        }

        // 等待可用带宽
        while self.used.load(Ordering::SeqCst) >= self.limit {
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    pub fn add(&self, n: u64) {
        self.used.fetch_add(n, Ordering::SeqCst);
    }
}

中继协议

// 中继消息格式
struct RelayMessage {
    source_id: [u8; 16],        // 源节点 ID (UUID)
    target_id: [u8; 16],        // 目标节点 ID (UUID)
    payload: Vec<u8>,           // 负载数据
    timestamp: u64,             // 时间戳
    sequence: u64,              // 序列号
}

// 加密
fn encrypt_payload(data: &[u8], key: &[u8]) -> Vec<u8> {
    // 使用 libsodium 加密
}

// 解密
fn decrypt_payload(data: &[u8], key: &[u8]) -> Vec<u8> {
    // 使用 libsodium 解密
}

连接管理

// src/peer.rs

pub struct ConnectionManager {
    // 在线节点 (ID -> ConnectionInfo)
    online_peers: Arc<DashMap<String, ConnectionInfo>>,

    // 地址映射 (SocketAddr -> ID)
    addr_to_id: Arc<DashMap<SocketAddr, String>>,

    // 消息通道
    msg_tx: mpsc::UnboundedSender<(String, Message)>,
}

pub struct ConnectionInfo {
    pub id: String,
    pub addr: SocketAddr,
    pub conn_type: ConnectionType,
    pub last_heartbeat: DateTime<Utc>,
    pub public_key: Vec<u8>,
}

impl ConnectionManager {
    /// 添加连接
    pub async fn add_connection(&self, id: String, addr: SocketAddr, conn_type: ConnectionType) {
        let info = ConnectionInfo {
            id: id.clone(),
            addr,
            conn_type,
            last_heartbeat: Utc::now(),
            public_key: Vec::new(),
        };

        self.online_peers.insert(id.clone(), info);
        self.addr_to_id.insert(addr, id);
    }

    /// 移除连接
    pub async fn remove_connection(&self, id: &str) {
        if let Some((addr, _)) = self.online_peers.remove(id) {
            self.addr_to_id.remove(&addr);
        }
    }

    /// 根据 ID 查找连接
    pub fn get_connection(&self, id: &str) -> Option<ConnectionInfo> {
        self.online_peers.get(id).map(|v| v.clone())
    }

    /// 更新心跳
    pub async fn update_heartbeat(&self, id: &str) {
        if let Some(mut info) = self.online_peers.get_mut(id) {
            info.last_heartbeat = Utc::now();
        }
    }

    /// 获取所有在线节点
    pub fn get_online_peers(&self) -> Vec<String> {
        self.online_peers.keys().map(|s| s.clone()).collect()
    }
}

用户认证 (可选)

// 用户认证模块

use jsonwebtoken::{decode, encode, Header, Validation};
use bcrypt::{hash, verify};

pub struct AuthManager {
    jwt_secret: Vec<u8>,
    jwt_issuer: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String,             // 用户 ID
    pub username: String,
    pub exp: usize,              // 过期时间
    pub iat: usize,              // 签发时间
}

impl AuthManager {
    /// 生成 JWT 令牌
    pub fn generate_token(&self, user_id: &str, username: &str) -> Result<String> {
        let exp = (Utc::now() + Duration::days(7)).timestamp() as usize;
        let iat = Utc::now().timestamp() as usize;

        let claims = Claims {
            sub: user_id.to_string(),
            username: username.to_string(),
            exp,
            iat,
        };

        encode(&Header::default(), &claims, &self.jwt_secret)
            .map_err(|e| Error::JwtError(e.to_string()))
    }

    /// 验证 JWT 令牌
    pub fn verify_token(&self, token: &str) -> Result<Claims> {
        let validation = Validation::default();
        decode::<Claims>(token, &self.jwt_secret, &validation)
            .map(|d| d.claims)
            .map_err(|e| Error::JwtError(e.to_string()))
    }

    /// 验证密码
    pub fn verify_password(&self, password: &str, hash: &str) -> bool {
        verify(password, hash).unwrap_or(false)
    }

    /// 哈希密码
    pub fn hash_password(&self, password: &str) -> Result<String> {
        hash(password, 12).map_err(|e| Error::BcryptError(e.to_string()))
    }
}

配置管理

# hbbs.toml 或命令行参数

[hbbs]
# 监听地址
host = "0.0.0.0"
port = 21116

# 数据库
db_file = "hbbs.db"

# 密钥
key = ""

# 中继服务器
relay = ""

# 用户认证
enable_user_auth = false

[relay]
# 中继服务器监听地址
host = "0.0.0.0"
port = 21117

# 带宽限制 (bytes/s)
max_bandwidth = 10485760  # 10 MB/s

# 最大连接数
max_connections = 1000

通信流程

┌─────────────┐                              ┌─────────────┐
│  Client A   │                              │  Client B   │
│  (ID: 1234) │                              │  (ID: 5678) │
└──────┬──────┘                              └──────┬──────┘
       │                                            │
       │  1. 连接 hbbs:21116                        │
       ├───────────────────────────────────────────►│
       │                                            │
       │  2. RegisterPeer (注册 ID)                 │
       │◄──────────────────────────────────────────┤
       │                                            │
       │  3. Client A 请求连接 B (PunchHoleRequest) │
       │◄──────────────────────────────────────────┤
       │                                            │
       │  4. hbbs 通知 B (PunchHole)                │
       │───────────────────────────────────────────►│
       │                                            │
       │  5a. P2P 成功: A 和 B 直接连接             │
       │◄──────────────────────────────────────────┤
       │                                            │
       │  5b. P2P 失败: 通过 hbbr 中继              │
       │◄══════ hbbs/hbbr 中继 ════════════════════│

可复用代码

模块 路径 用途
rendezvous_server src/rendezvous_server.rs 信令服务器核心
relay_server src/relay_server.rs 中继服务器
peer src/peer.rs 节点管理
database src/database.rs SQLite 操作

在本项目中的使用

作为 data-server,提供:

┌─────────────────────────────────────────────────────────────────┐
│                        data-server (rustdesk-server)              │
│                                                                 │
│  ┌─────────────────────┐    ┌─────────────────────┐            │
│  │      hbbs           │    │       hbbr          │            │
│  │  (信令服务器)       │    │   (中继服务器)       │            │
│  │  - 端口: 21116      │    │   - 端口: 21117     │            │
│  │  - ID 注册         │    │   - 数据中继        │            │
│  │  - 连接协商        │    │   - 带宽限制        │            │
│  │  - NAT 穿透        │    │   - TCP/WS          │            │
│  └─────────────────────┘    └─────────────────────┘            │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                     agent-client / agent-server-admin            │
└─────────────────────────────────────────────────────────────────┘

与 qiming-rustdesk 的配合

qiming-rustdesk (客户端) <───────────────► rustdesk-server (data-server)
    │                                              │
    │  连接 hbbs:21116                             │
    │  ├── RegisterPeer (注册 ID)                  │
    │  ├── PunchHoleRequest (请求连接)             │
    │  └── RequestRelay (请求中继)                 │
    │                                              │
    │  连接 hbbr:21117 (中继模式)                  │
    │  └── 数据中继传输                            │

在本项目中的使用场景

场景1agent-server-admin 连接 agent-client

// rustdesk-server 作为消息中转枢纽

// 1. agent-client 连接 hbbs注册 ID
let client_id = register_peer_to_hbbs(client).await?;

// 2. agent-server-admin 连接 hbbs查询目标客户端
let target_info = query_peer_from_hbbs(target_id).await?;

// 3. hbbs 协调双方建立 P2P 连接或中继
let connection = establish_connection(source_id, target_id).await?;

// 4. 双方通过建立的通道传输 protobuf 消息
send_protobuf_message(&connection, message).await?;

场景2SSE 实时消息推送

// hbbs 支持 WebSocket 连接,用于 SSE 推送

use tokio_tungstenite::{WebSocketStream, tungstenite::Message};
use futures::{StreamExt, SinkExt};

pub struct SseHandler {
    // 客户端连接 (client_id -> WebSocket)
    connections: Arc<DashMap<String, WebSocketStream<TcpStream>>>,
}

impl SseHandler {
    /// 处理 SSE 连接
    pub async fn handle_sse_connection(
        &self,
        client_id: String,
        ws_stream: WebSocketStream<TcpStream>,
    ) -> Result<()> {
        let (mut ws_sender, mut ws_receiver) = ws_stream.split();

        // 保存连接
        self.connections.insert(client_id.clone(), ws_stream);

        // 接收消息
        while let Some(msg) = ws_receiver.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    self.handle_client_message(&client_id, &text).await?;
                }
                Ok(Message::Close(_)) => {
                    self.connections.remove(&client_id);
                    break;
                }
                _ => {}
            }
        }

        Ok(())
    }

    /// 向客户端推送消息
    pub async fn push_message(&self, client_id: &str, message: &str) -> Result<()> {
        if let Some((_, ws_stream)) = self.connections.get(client_id) {
            let mut ws = ws_stream.lock().await;
            ws.send(Message::Text(message.to_string())).await?;
        }
        Ok(())
    }

    /// 广播消息给所有在线客户端
    pub async fn broadcast(&self, message: &str) {
        for (client_id, _) in self.connections.iter() {
            let _ = self.push_message(client_id, message).await;
        }
    }
}

场景3文件传输中继

// hbbr 支持大文件分片中继传输

pub struct FileRelayHandler {
    // 文件传输会话 (session_id -> FileSession)
    sessions: Arc<DashMap<String, FileSession>>,
    // 带宽限制
    speed_limit: Arc<SpeedLimit>,
}

pub struct FileSession {
    pub source_id: String,
    pub target_id: String,
    pub file_name: String,
    pub file_size: u64,
    pub chunks: u64,          // 总块数
    pub current_chunk: u64,   // 当前块
    pub last_update: DateTime<Utc>,
}

impl FileRelayHandler {
    /// 开始文件传输会话
    pub async fn start_session(
        &self,
        source_id: &str,
        target_id: &str,
        file_name: &str,
        file_size: u64,
    ) -> Result<String> {
        let chunk_size = 64 * 1024; // 64KB
        let chunks = (file_size + chunk_size - 1) / chunk_size;

        let session_id = uuid::Uuid::new_v4().to_string();
        let session = FileSession {
            source_id: source_id.to_string(),
            target_id: target_id.to_string(),
            file_name: file_name.to_string(),
            file_size,
            chunks,
            current_chunk: 0,
            last_update: Utc::now(),
        };

        self.sessions.insert(session_id.clone(), session);
        Ok(session_id)
    }

    /// 中继文件块
    pub async fn relay_chunk(
        &self,
        session_id: &str,
        chunk_data: &[u8],
        offset: u64,
    ) -> Result<()> {
        // 应用带宽限制
        self.speed_limit.wait_if_needed().await;

        if let Some(session) = self.sessions.get(session_id) {
            // 找到目标连接并转发
            if let Some(target_conn) = self.get_connection(&session.target_id).await {
                let relay_msg = FileChunkMessage {
                    session_id: session_id.to_string(),
                    offset,
                    data: chunk_data.to_vec(),
                    timestamp: Utc::now().timestamp_millis(),
                };
                self.forward_to(&target_conn, &relay_msg).await?;
            }

            // 更新进度
            self.sessions.insert(session_id.to_string(), FileSession {
                current_chunk: offset / 64_1024 + 1,
                ..session.value().clone()
            });
        }

        Ok(())
    }

    /// 清理超时会话
    pub async fn cleanup_timeout_sessions(&self, timeout: Duration) {
        let now = Utc::now();
        let expired: Vec<String> = self.sessions
            .iter()
            .filter(|s| (now - s.last_update) > timeout)
            .map(|s| s.key().clone())
            .collect();

        for session_id in expired {
            self.sessions.remove(&session_id);
        }
    }
}

场景4心跳保活机制

// hbbs 维护客户端在线状态,定期发送心跳

pub struct HeartbeatManager {
    // 客户端心跳 (client_id -> last_heartbeat)
    heartbeats: Arc<DashMap<String, DateTime<Utc>>>,

    // 心跳间隔
    heartbeat_interval: Duration,

    // 超时时间
    timeout: Duration,

    // 清理任务
    cleanup_task: JoinHandle<()>,
}

impl HeartbeatManager {
    pub fn new(heartbeat_interval: Duration, timeout: Duration) -> Self {
        let heartbeats = Arc::new(DashMap::new());
        let this = Self {
            heartbeats,
            heartbeat_interval,
            timeout,
            cleanup_task: tokio::spawn(Self::cleanup_loop(timeout)),
        };
        this
    }

    /// 更新心跳
    pub fn update_heartbeat(&self, client_id: &str) {
        self.heartbeats.insert(client_id.to_string(), Utc::now());
    }

    /// 检查客户端是否在线
    pub fn is_online(&self, client_id: &str) -> bool {
        if let Some(heartbeat) = self.heartbeats.get(client_id) {
            (Utc::now() - *heartbeat) < self.timeout
        } else {
            false
        }
    }

    /// 获取所有在线客户端
    pub fn get_online_clients(&self) -> Vec<String> {
        let now = Utc::now();
        self.heartbeats
            .iter()
            .filter(|(_, hb)| (now - *hb) < self.timeout)
            .map(|s| s.key().clone())
            .collect()
    }

    async fn cleanup_loop(timeout: Duration) {
        let mut interval = tokio::time::interval(Duration::from_secs(60));
        loop {
            interval.tick().await;
            // 清理超时的客户端心跳记录
            // 这是一个简化实现,实际应该使用 ArcSwap 或其他原子操作
        }
    }
}

场景5连接状态统计

// 统计客户端连接状态(用于 UI 显示)

pub struct ConnectionStats {
    // 总连接数
    total_connections: AtomicUsize,

    // P2P 连接数
    p2p_connections: AtomicUsize,

    // 中继连接数
    relay_connections: AtomicUsize,

    // 按 NAT 类型统计
    nat_type_stats: HashMap<String, AtomicUsize>,

    // 在线客户端列表
    online_clients: DashMap<String, ClientInfo>,
}

#[derive(Clone)]
pub struct ClientInfo {
    pub id: String,
    pub conn_type: ConnectionType,
    pub nat_type: String,
    pub connected_at: DateTime<Utc>,
    pub last_activity: DateTime<Utc>,
}

impl ConnectionStats {
    /// 记录新连接
    pub fn record_connection(&self, client_id: &str, conn_type: ConnectionType, nat_type: &str) {
        self.total_connections.fetch_add(1, Ordering::SeqCst);

        match conn_type {
            ConnectionType::P2P => self.p2p_connections.fetch_add(1, Ordering::SeqCst),
            ConnectionType::Relay => self.relay_connections.fetch_add(1, Ordering::SeqCst),
            _ => {}
        }

        let stats = self.nat_type_stats.entry(nat_type.to_string()).or_insert_with(|| AtomicUsize::new(0));
        stats.fetch_add(1, Ordering::SeqCst);

        self.online_clients.insert(client_id.to_string(), ClientInfo {
            id: client_id.to_string(),
            conn_type,
            nat_type: nat_type.to_string(),
            connected_at: Utc::now(),
            last_activity: Utc::now(),
        });
    }

    /// 获取统计摘要
    pub fn get_summary(&self) -> ConnectionSummary {
        ConnectionSummary {
            total_online: self.online_clients.len(),
            p2p_count: self.p2p_connections.load(Ordering::SeqCst),
            relay_count: self.relay_connections.load(Ordering::SeqCst),
            nat_distribution: self.nat_type_stats
                .iter()
                .map(|(k, v)| (k.clone(), v.load(Ordering::SeqCst)))
                .collect(),
        }
    }
}

关键设计模式

  1. 内存缓存 + SQLite 持久化: PeerMap 使用 RwLock 内存缓存 + SQLite 持久化
  2. DashMap 并发管理: ConnectionManager 使用 DashMap 管理在线节点
  3. 心跳检测: 定期检测连接活性,清理超时连接
  4. 带宽限制: 中继服务器使用滑动窗口进行带宽控制