charging-cabinet/software/device-server/src/tcp/server.rs

//! TCP 服务端：监听端口、接收设备连接、按 LF 分割 JSON 消息。

use futures::StreamExt;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tokio_util::codec::{FramedRead, LinesCodec};

use super::connection::{ConnectionPool, DeviceConnection};
use super::handler::handle_message;
use super::protocol::DeviceMessage;

/// 未登录连接**允许**发送的唯一两个 `act`（协议定稿：其余指令必须先 `login` 成功）。
const UNAUTH_ALLOWED_ACTS: &[&str] = &["auth_str", "login"];

/// 连接建立后必须在此秒数内完成 `login` 认证，否则平台主动断开。
/// 依据：协议规定"login 成功这个 TCP 才可以保持"。
const UNAUTH_TIMEOUT_SECS: u64 = 30;

/// 未登录超时的巡检间隔（秒）。
const UNAUTH_CHECK_INTERVAL_SECS: u64 = 5;

/// 启动 TCP 监听循环。
pub async fn run_tcp_server(
    addr: String,
    pool: ConnectionPool,
    redis: redis::aio::ConnectionManager,
    node_id: String,
) {
    let listener = TcpListener::bind(&addr)
        .await
        .expect("TCP 端口绑定失败");
    tracing::info!("TCP 设备服务启动: {}", addr);

    loop {
        match listener.accept().await {
            Ok((stream, peer)) => {
                tracing::info!("[新连接] peer={}", peer);
                // ponytail: 不需要设 TCP keepalive，我们是 server，设备才是 client。
                // 设备每 60s 发 status_post 就是应用层保活。
                let pool = pool.clone();
                let redis = redis.clone();
                let nid = node_id.clone();
                tokio::spawn(handle_connection(stream, peer, pool, redis, nid));
            }
            Err(e) => {
                tracing::error!("accept 失败: {}", e);
            }
        }
    }
}

/// 处理单条 TCP 连接。
async fn handle_connection(
    stream: TcpStream,
    peer: std::net::SocketAddr,
    pool: ConnectionPool,
    redis: redis::aio::ConnectionManager,
    node_id: String,
) {
    let (reader, writer) = tokio::io::split(stream);
    // 单行上限 64 KB：status_post 为「每 60 秒全量上报」
    // （当前 6板×6组 ≈ 8 KB；设计上限 15板×8组 ≈ 25 KB，含电芯数据约 40 KB），留足余量
    let mut lines = FramedRead::new(reader, LinesCodec::new_with_max_length(65536));

    // 写通道：HTTP 命令下发通过此通道写入 TCP，缓冲区 100 条防堆积
    let (tx, mut rx) = mpsc::channel::<String>(100);

    let mut current_dev_id: Option<String> = None;
    let mut writer = writer;
    let connected_at = std::time::Instant::now();

    loop {
        tokio::select! {
            // 读取设备消息
            line = lines.next() => {
                match line {
                    Some(Ok(text)) => {
                        if text.trim().is_empty() {
                            continue;
                        }

                        let msg: DeviceMessage = match serde_json::from_str(&text) {
                            Ok(m) => m,
                            Err(e) => {
                                tracing::warn!("[{}] JSON 解析失败: {} raw={}", peer, e, text);
                                continue;
                            }
                        };

                        // 如果是 login/auth_str，需要 dev_id 来关联连接
                        // 注意：注册在签名验证之前，因为 login 响应（auth_str/cell）必须通过
                        // 连接池下发。先注册才能回复。纯内网 TCP + 真实设备重连快，攻击面有限。
                        // 审查编号：I3（已确认不修，协议约束）
                        if (msg.act == "auth_str" || msg.act == "login")
                            && let Some(ref dev_id) = msg.dev_id
                                && current_dev_id.is_none() {
                                    let conn = DeviceConnection::new(tx.clone());
                                    pool.register(dev_id.clone(), conn).await;
                                    current_dev_id = Some(dev_id.clone());
                                    // 首次识别设备时写入 route
                                    update_device_route(&redis, dev_id, &node_id).await;
                                }

                        // 更新活跃时间，同时刷新 route（处理重连切换到其他节点的情况）
                        if let Some(ref dev_id) = current_dev_id {
                            pool.touch(dev_id).await;
                            update_device_route(&redis, dev_id, &node_id).await;
                        }

                        // 安全规则（协议定稿）：除 auth_str / login 外，其余指令必须先 login 成功
                        // 未登录连接发来的业务报文一律丢弃（连接尚未与 dev_id 建立可信绑定，
                        // 平台侧也无法路由应答）
                        if !UNAUTH_ALLOWED_ACTS.contains(&msg.act.as_str()) {
                            let authed = match current_dev_id.as_deref() {
                                Some(dev_id) => pool.is_authenticated(dev_id).await,
                                None => false,
                            };
                            if !authed {
                                tracing::warn!(
                                    "[安全] 连接未登录，丢弃报文 act={} peer={}",
                                    msg.act,
                                    peer
                                );
                                continue;
                            }
                        }

                        // 处理消息（传入连接绑定的设备身份：新协议业务上行不带 dev_id，
                        // 转发到 Redis 前由 handler 补写，api-server 依赖它识别设备）。
                        // 无返回值：本服务不直接回复设备，应答由 api-server 经 device:commands 下发。
                        handle_message(msg, current_dev_id.as_deref(), &pool, &redis).await;
                    }
                    Some(Err(e)) => {
                        tracing::warn!("[{}] 读取错误: {}", peer, e);
                        break;
                    }
                    None => {
                        tracing::info!("[{}] 连接关闭", peer);
                        break;
                    }
                }
            }
            // 未登录超时巡检：连接建立后 UNAUTH_TIMEOUT_SECS 秒内未通过 login 即断开
            // （协议规定：login 成功这条 TCP 才可以保持）
            _ = tokio::time::sleep(std::time::Duration::from_secs(UNAUTH_CHECK_INTERVAL_SECS)) => {
                use tokio::io::AsyncWriteExt;

                let authed = match current_dev_id.as_deref() {
                    Some(dev_id) => pool.is_authenticated(dev_id).await,
                    None => false,
                };
                if !authed && connected_at.elapsed().as_secs() >= UNAUTH_TIMEOUT_SECS {
                    tracing::warn!(
                        "[{}] 未在 {} 秒内完成 login，平台主动断开连接",
                        peer,
                        UNAUTH_TIMEOUT_SECS
                    );
                    let _ = writer.shutdown().await;
                    break;
                }
            }
            // 向设备写指令
            Some(json) = rx.recv() => {
                use tokio::io::AsyncWriteExt;

                // api-server 会附加内部标记 `_close`（login 失败后断开）与
                // `_dev_id`（路由目标）：两者仅供服务端使用，下发前摘除，设备侧不可见
                let (payload, close_after) = split_internal_flags(&json);

                let framed = format!("{}\n", payload);
                if let Err(e) = writer.write_all(framed.as_bytes()).await {
                    tracing::error!("[{}] 写入失败: {}", peer, e);
                    break;
                }

                // 规范：login 失败即断开 TCP，强制设备重连并重新走握手流程
                if close_after {
                    tracing::warn!(
                        "[{}] login 失败，主动断开连接 dev_id={:?}",
                        peer,
                        current_dev_id
                    );
                    let _ = writer.shutdown().await;
                    break;
                }
            }
        }
    }

    // 连接断开，清理
    if let Some(ref dev_id) = current_dev_id {
        pool.remove(dev_id).await;
        tracing::info!("[断连] dev_id={}", dev_id);
    }
}

/// 更新设备路由：在 `device:{imei}` hash 中写入 `node_id`，并刷新 TTL。
/// 路由信息用于多节点时指令消费方判断设备连接在哪个节点上。
async fn update_device_route(
    redis: &redis::aio::ConnectionManager,
    dev_id: &str,
    node_id: &str,
) {
    let route_key = format!("device:{}", dev_id);
    let mut conn = redis.clone();
    let _: Result<(), _> = redis::cmd("HSET")
        .arg(&route_key)
        .arg("node_id")
        .arg(node_id)
        .query_async(&mut conn)
        .await;
    // ponytail: 与设备在线 TTL 保持一致（300s），HSET 不支持同时设置 TTL
    let _: Result<(), _> = redis::cmd("EXPIRE")
        .arg(&route_key)
        .arg(300_u32)
        .query_async(&mut conn)
        .await;
}

/// 摘除 api-server 附加的**内部标记**（`_close` / `_dev_id`）。
///
/// 返回 `(实际下发给设备的报文, 是否在发送后断开连接)`。
/// 这两个字段只用于服务端内部：
/// - `_close`：`login` 失败后主动断开连接
/// - `_dev_id`：指令的路由目标（设备身份），下发前必须摘除 ——
///   否则设备会把它当成协议字段，且协议定稿后 `dev_id` 只在 `auth_str`/`login` 携带
///
/// 两者**都不会下发到设备**。
fn split_internal_flags(json: &str) -> (String, bool) {
    match serde_json::from_str::<serde_json::Value>(json) {
        Ok(mut value) => {
            let close_after = value
                .get("_close")
                .and_then(|flag| flag.as_bool())
                .unwrap_or(false);
            if let Some(obj) = value.as_object_mut() {
                obj.remove("_close");
                obj.remove("_dev_id");
            }
            (value.to_string(), close_after)
        }
        // 非 JSON（理论上不会出现）：原样下发，不断开
        Err(_) => (json.to_string(), false),
    }
}

// ---------------------------------------------------------------------------
// 单元测试：内部字段摘除（`_close` / `_dev_id` 绝不能下发到设备）
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
    use super::split_internal_flags;

    /// 普通指令：无内部字段时原样返回、不断开
    #[test]
    fn test_split_internal_flags_plain_command_untouched() {
        let raw = r#"{"act":"on","msg_id":1,"id":"01-3"}"#;
        let (payload, close) = split_internal_flags(raw);
        assert!(!close);
        let v: serde_json::Value = serde_json::from_str(&payload).unwrap();
        assert_eq!(v["act"], "on");
        assert_eq!(v["msg_id"], 1);
        assert_eq!(v["id"], "01-3");
        assert!(v.get("_dev_id").is_none());
        assert!(v.get("_close").is_none());
    }

    /// `_dev_id`（路由目标）必须被摘除 —— 它是服务端内部字段，
    /// 且协议定稿后 `dev_id` 只在 auth_str/login 携带，泄漏会让设备误解报文语义
    #[test]
    fn test_split_internal_flags_removes_route_dev_id() {
        let raw = r#"{"act":"open","msg_id":2,"_dev_id":"864720562557021"}"#;
        let (payload, close) = split_internal_flags(raw);
        assert!(!close);
        let v: serde_json::Value = serde_json::from_str(&payload).unwrap();
        assert_eq!(v["act"], "open");
        assert_eq!(v["msg_id"], 2);
        assert!(v.get("_dev_id").is_none());
    }

    /// `_close` 被摘除并回报"发送后断开"（login 失败即断连）
    #[test]
    fn test_split_internal_flags_reports_and_removes_close() {
        let raw = r#"{"act":"login","msg_id":2,"suc":0,"err":"invalid signature","_close":true}"#;
        let (payload, close) = split_internal_flags(raw);
        assert!(close);
        let v: serde_json::Value = serde_json::from_str(&payload).unwrap();
        assert_eq!(v["act"], "login");
        assert_eq!(v["err"], "invalid signature");
        assert!(v.get("_close").is_none());
    }

    /// 两个内部字段同时出现时都要摘掉
    #[test]
    fn test_split_internal_flags_removes_both() {
        let raw = r#"{"act":"login","suc":0,"_dev_id":"864720562557021","_close":true}"#;
        let (payload, close) = split_internal_flags(raw);
        assert!(close);
        let v: serde_json::Value = serde_json::from_str(&payload).unwrap();
        assert!(v.get("_dev_id").is_none());
        assert!(v.get("_close").is_none());
    }

    /// 非 JSON 原样下发且不断开（理论上不会出现，但不能因此丢失报文）
    #[test]
    fn test_split_internal_flags_non_json_passthrough() {
        let (payload, close) = split_internal_flags("not-json");
        assert_eq!(payload, "not-json");
        assert!(!close);
    }
}