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, RegisterOutcome, next_conn_id};
use super::connection::{clear_device_authed, refresh_device_authed};
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);

    // **主动驱逐信号**：本连接的池条目一旦离开池（被新连接接管），条目里的
    // `DeviceConnection::close`（oneshot 发送端）随之 drop → 下面 `select!` 的接收端
    // 立刻就绪 → 本连接立即关闭 socket 并退出。不依赖设备下次上行，
    // 也**不新增任何协议报文**（协议冻结，不能用自造帧通知设备）。
    let (close_tx, mut close_rx) = tokio::sync::oneshot::channel::<()>();

    let mut current_dev_id: Option<String> = None;
    let mut writer = writer;
    let connected_at = std::time::Instant::now();
    // 本 TCP 连接的唯一标识：池条目的创建/删除/认证态判定都以它校验归属。
    // `dev_id` 来自报文（不可信），只有 `conn_id` 能证明"池里那条entry是不是我的"。
    let conn_id = next_conn_id();

    // 连接状态在首次 `auth_str`/`login` 注册时才进池；这里先构造好放在 `Option` 里，
    // 注册时 `take()`（`close_tx` 只能移动一次，而注册代码在循环体内）。
    // `peer` 一并带进池：下行拦截 / 连接接管类 WARN 需要定位到具体 socket。
    let mut pending_conn = Some(DeviceConnection::new(tx.clone(), conn_id, peer, close_tx));

    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 来关联连接
                        //
                        // ⚠️ 安全：注册发生在**签名校验之前**（`auth_str`/`login` 的应答
                        // 必须经连接池下发，先注册才能回复），所以报文里的 `dev_id` 是
                        // **不可信输入**。原注释曾以"纯内网 TCP + 真实设备重连快"为由
                        // 认为攻击面有限（审查编号 I3，不修）—— 该前提**不成立**：
                        // 9001 必须对 4G 设备公网可达，任何能连上端口的人都能伪造 IMEI。
                        // 真实威胁（C1）：
                        //   `nc host 9001` 发 `{"act":"auth_str","dev_id":"<受害IMEI>","msg_id":1}`
                        //   1) 覆盖式 register 会把真机的池条目换成攻击者连接 →
                        //      真机此后所有上行被判"未登录"丢弃；
                        //   2) 攻击者断开时无条件 remove 又删掉真机条目 →
                        //      平台永久"设备不在线"，每 25 秒重连一次即可维持永久 DoS。
                        // 本次修法（不改 TCP 协议，纯服务端逻辑）：
                        //   1) 每条连接分配唯一 `conn_id`，池操作按 conn_id 校验归属；
                        //   2) `register` **允许接管**同 dev_id 的旧条目（未认证占位 / 已认证连接，
                        //      MQTT 同 clientId 接管语义），接管时显式调用 `evict()` 主动驱逐旧连接
                        //      —— 既不拒绝真机重连（4G 静默掉线后旧 socket 可能仍是 ESTABLISHED），
                        //      也不会让攻击者把真机挤下线后就无人能接管回来；
                        //   3) 断连清理改用 `remove_if(dev_id, conn_id)`，只删自己的条目；
                        //   4) 认证态与活跃时间都按 `conn_id` 归属判定，被接管的连接不能借用
                        //      新连接的已认证身份、也不能替它刷新 `last_seen`。
                        if (msg.act == "auth_str" || msg.act == "login")
                            && let Some(ref dev_id) = msg.dev_id
                                && current_dev_id.is_none()
                                    && let Some(conn) = pending_conn.take() {
                                    match pool.register(dev_id.clone(), conn).await {
                                        RegisterOutcome::Inserted => {
                                            current_dev_id = Some(dev_id.clone());
                                            // 首次识别设备时写入 route
                                            update_device_route(&redis, dev_id, &node_id).await;
                                        }
                                        RegisterOutcome::ReplacedUnauthenticated { evicted_conn_id } => {
                                            // 顶替未认证占位（攻击者占位 / 真机重连而旧连接未被清理）
                                            tracing::warn!(
                                                "[安全] dev_id={} 的未认证占位连接(conn_id={}) 已被本连接(conn_id={}, peer={}) 顶替并驱逐",
                                                dev_id,
                                                evicted_conn_id,
                                                conn_id,
                                                peer
                                            );
                                            current_dev_id = Some(dev_id.clone());
                                            update_device_route(&redis, dev_id, &node_id).await;
                                            // 连接级认证标记必须**立即清除**：此刻持有该 IMEI 的是
                                            // 一条**尚未认证**的新连接，api-server 不能再认为它"可下发"。
                                            // 新连接 `login` 成功后标记会重新写入。
                                            clear_device_authed(&redis, dev_id).await;
                                        }
                                        RegisterOutcome::ReplacedAuthenticated { evicted_conn_id } => {
                                            // **连接接管**：与 MQTT 同 clientId 语义一致（新连接赢）。
                                            // 典型场景是真机静默掉线（NAT 超时 / 4G 信号丢失）后重连，
                                            // 而服务端侧旧 socket 仍被视为 ESTABLISHED。若这里拒绝新连接，
                                            // 真机发什么都进不了业务链路（可用性事故），故必须接管。
                                            // 旧连接由「主动驱逐句柄」立即断开（见 handle_connection 开头的
                                            // close_rx 分支），不需要设备发任何报文配合。
                                            // 这条 WARN 是排查"设备反复掉线/被顶替"的关键信号，务必保留。
                                            tracing::warn!(
                                                "[安全] 同一 IMEI 发生连接接管：dev_id={} 的已认证连接(conn_id={}) 已被新连接(conn_id={}, peer={}) 接管并驱逐（新连接需重新 login 认证）",
                                                dev_id,
                                                evicted_conn_id,
                                                conn_id,
                                                peer
                                            );
                                            current_dev_id = Some(dev_id.clone());
                                            update_device_route(&redis, dev_id, &node_id).await;
                                            // 同上：旧连接的认证态不能顺延给新连接（接管方须重新 login）
                                            clear_device_authed(&redis, dev_id).await;
                                        }
                                    }
                                }

                        // 接管兜底（安全网）：池条目若已不属于本连接，说明本连接已被新连接接管
                        // （等价于"superseded 标记"）。主动驱逐信号（close_rx）通常已触发，
                        // 这里覆盖"信号尚未被 select! 处理完"的竞态：立即断开，不再处理任何上行。
                        if let Some(ref dev_id) = current_dev_id
                            && !pool.is_owner(dev_id, conn_id).await {
                            use tokio::io::AsyncWriteExt;
                            tracing::warn!(
                                "[安全] 连接(conn_id={}, dev_id={}, peer={}) 已被新连接接管，平台主动断开且不再处理其上行",
                                conn_id,
                                dev_id,
                                peer
                            );
                            let _ = writer.shutdown().await;
                            break;
                        }

                        // 更新活跃时间，同时刷新 route（处理重连切换到其他节点的情况）。
                        // `touch` 带 conn_id：只有当前持有者能刷新活跃时间（被接管的旧连接不行）。
                        if let Some(ref dev_id) = current_dev_id {
                            pool.touch(dev_id, conn_id).await;
                            update_device_route(&redis, dev_id, &node_id).await;
                        }

                        // 安全规则（协议定稿）：除 auth_str / login 外，其余指令必须先 login 成功
                        // 未登录连接发来的业务报文一律丢弃（连接尚未与 dev_id 建立可信绑定，
                        // 平台侧也无法路由应答）
                        //
                        // 认证态按 `conn_id` 归属判定：占位连接即使被真机顶替、且真机已认证，
                        // 也不能"借用"到已认证身份伪造该 IMEI 的上行。
                        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, conn_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;

                        // 连接级认证标记：已登录连接**每次上行**刷新 TTL（`device:authed:{imei}`），
                        // 供 api-server 判断"该 IMEI 是否存在已登录连接"（协议 §4.2）。
                        // 这里用 `EXPIRE`（只刷新已存在的键）—— 未登录连接不可能凭空造出该标记。
                        if let Some(ref dev_id) = current_dev_id
                            && pool.is_authenticated(dev_id, conn_id).await
                        {
                            refresh_device_authed(&redis, dev_id).await;
                        }
                    }
                    Some(Err(e)) => {
                        tracing::warn!("[{}] 读取错误: {}", peer, e);
                        break;
                    }
                    None => {
                        tracing::info!("[{}] 连接关闭", peer);
                        break;
                    }
                }
            }
            // 主动驱逐：本连接的池条目已被新连接接管（连接池调用 `evict()` 发来信号）。
            // 立即断开，让接管方独占该 dev_id；设备侧无感知（不涉及任何协议报文），
            // 真机在自己的重连流程里会重新 auth_str/login。
            _ = &mut close_rx => {
                use tokio::io::AsyncWriteExt;

                tracing::warn!(
                    "[安全] 连接(conn_id={}, dev_id={:?}, peer={}) 的池条目已被新连接接管，平台主动断开旧连接",
                    conn_id,
                    current_dev_id,
                    peer
                );
                let _ = writer.shutdown().await;
                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, conn_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;
                }
            }
        }
    }

    // 连接断开，清理：**只删 conn_id 属于自己的那条条目**。
    // 无条件 `remove(dev_id)` 会让攻击者用受害 IMEI 占位后、在断开时把真机的
    // 池条目一并删掉 → 平台永久判定该设备离线（C1 的 DoS 主路径）。
    if let Some(ref dev_id) = current_dev_id {
        if pool.remove_if(dev_id, conn_id).await {
            tracing::info!("[断连] dev_id={} conn_id={} 池条目已清理", dev_id, conn_id);
            // 本连接是该 IMEI 的**当前持有者**且已断开 ⇒ 连接级认证标记必须立刻删除，
            // 否则 api-server 在 TTL 窗口内仍会认为"设备在线可下发"（本标记要消灭的正是这类假在线）。
            clear_device_authed(&redis, dev_id).await;
        } else {
            tracing::info!(
                "[断连] dev_id={} conn_id={} 池条目已易主，保留当前持有者（不清理）",
                dev_id,
                conn_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);
    }
}