futu_server/

conn.rs

// 单连接管理：状态机、帧收发、加密、心跳超时

use std::collections::HashSet;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Instant;

use bytes::Bytes;
use futures::{SinkExt, StreamExt};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio_util::codec::Framed;

use futu_auth::Scope;
use futu_codec::FutuCodec;
use futu_codec::frame::FutuFrame;
use futu_codec::header::ProtoFmtType;
use futu_core::error::FutuError;
use futu_net::encrypt;

/// 连接状态
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum ConnState {
    /// TCP / WebSocket 刚建立，尚未完成 InitConnect 握手
    Connected,
    /// 已完成 InitConnect，等待首次业务请求
    Initialized,
    /// 正常交互中（业务请求 / KeepAlive / push 均已流转）
    Active,
    /// 连接已断开（客户端主动关闭 / 被动超时 / IO 错误）
    Disconnected,
}

/// 单个客户端连接
pub struct ClientConn {
    /// 随机连接 ID（对应 C++ `GetRand_MilliTimeAndU22`）
    pub conn_id: u64,
    /// 连接状态：InitConnect 前 / 后 / 已断开
    pub state: ConnState,
    /// 随机 AES-128 key（InitConnect 响应里下发给客户端）
    pub aes_key: [u8; 16],
    /// AES 加解密已启用（InitConnect 完成且配置了 RSA 时为 true）
    pub aes_encrypt_enabled: bool,
    /// 该连接协商的 proto 格式（Protobuf / JSON）
    pub proto_fmt_type: ProtoFmtType,
    /// 上次收到 KeepAlive 的时间，用于超时检查
    pub last_keepalive: Instant,
    /// InitConnect.C2S.recvNotify：此连接是否接收市场状态 / 交易解锁等通知。
    ///
    /// C++ 在 `APIServer_InitConnect.cpp` 里把该字段写入 ConnInfo；
    /// `RegQotPush` / `Qot_Sub(isRegOrUnRegPush)` 不会修改这个开关。
    pub recv_notify: bool,
    /// 已收到的 KeepAlive 计数（监控用）
    pub keepalive_count: AtomicU32,
    /// 发送帧到此连接
    pub tx: mpsc::Sender<FutuFrame>,

    // ---- v1.0 WS per-message scope 鉴权（raw TCP legacy 兼容：scopes 空集全放行）----
    /// 该连接绑定的 API key id；WS 握手时填，未配 keys.json 时为 None
    pub key_id: Option<String>,
    /// 该连接持有的 scope 集合；空集 = legacy 模式 / TCP 直连，scope 检查放行
    pub scopes: HashSet<Scope>,
    /// v1.4.105 D3 (Phase 4) T-B2: 该连接 caller key 的 `allowed_markets` 硬
    /// 限额 (大写字符串 set, e.g. {"HK","US"}). `None` = 无限制 (legacy 模式
    /// / TCP 直连默认全开 / 未配 allowed_markets); `Some(set)` 非空 → push 端
    /// 应过滤 trd_market 不在 set 中的 trade event.
    ///
    /// **触发**: WS handshake 时从 `KeyRecord.allowed_markets` 拷贝过来.
    /// `PushDispatcher::push_trd_acc` 端 Layer 3 filter 检查. 与
    /// `caller_allowed_acc_ids` (Layer 1, per-call snapshot in IncomingRequest)
    /// 区别: 本字段是 per-conn snapshot (handshake 时一次性), 不随 per-call
    /// 重读 — KeyRecord SIGHUP reload 后**仅新建连接生效**, 老连接保持 snapshot
    /// (与 `scopes` / `caller_allowed_acc_ids` 的 snapshot 语义一致).
    pub allowed_markets: Option<std::sync::Arc<HashSet<String>>>,
    /// codex round 1 F4 (P2) v1.4.105: 该连接 caller key 的 `allowed_acc_ids`
    /// 硬限额 (per-conn snapshot, handshake 时一次性). `None` / `Some(empty)` =
    /// 无限制 (legacy 模式 / TCP 直连默认全开 / 未配 allowed_acc_ids);
    /// `Some(non-empty set)` →
    /// `PushDispatcher::push_trd_acc` 端 push-time 硬过滤 acc_id 不在 set 中
    /// 的 trade event (Layer 1, 与 `allowed_markets` 的 Layer 3 互补).
    /// Deny-all 使用 sentinel `{0}`，不使用空集合。
    ///
    /// **触发**: codex F4 指出 raw TCP push 端只查 `acc:read` scope +
    /// `allowed_markets`, 不查 `allowed_acc_ids`. 即使 request-time
    /// `SubAccPushHandler` 已阻止越权订阅, stale subscription / KeyRecord
    /// reload 后窄化的 acc 范围 / 历史 bug 留下的 conn→acc 关系 仍可能让 push
    /// 漏 leak. 本字段提供第二层 push-time 兜底.
    ///
    /// 与 `caller_allowed_acc_ids` (IncomingRequest, per-call) 区别: 本字段
    /// 在 push-time 用 (无 IncomingRequest), per-conn snapshot 与 `scopes` /
    /// `allowed_markets` 的 snapshot 语义一致.
    pub allowed_acc_ids: Option<std::sync::Arc<HashSet<u64>>>,
}

/// 从连接接收到的请求
#[derive(Debug)]
pub struct IncomingRequest {
    /// 发送请求的连接 ID（用于响应路由 + SubscriptionManager / cache 记账）
    ///
    /// **跨 surface 命名空间分配**（v1.4.106 codex 0517 ζ25-redo F2 沉淀）：
    /// - raw TCP listener: `ClientConn::generate_conn_id()` 派生（u32 范围）
    /// - REST: `crates/futu-rest/src/routes/qot.rs::REST_SHARED_CONN`
    ///   = `0xFFFF_FFFE`（u32 上限附近, 单值共享）
    /// - gRPC: `crates/futu-grpc/src/auth.rs::GRPC_STABLE_CONN_NAMESPACE`
    ///   = `0x4000_0000_0000_0000`（bit 62 namespace, 按 caller 派生）
    /// - WS / MCP: 通常派生自所属物理 TCP 连接的 conn_id
    ///
    /// 各 surface 不重叠. 加新 surface 时分配一个 namespace base, 不要与
    /// 上述 4 个段重合.
    pub conn_id: u64,
    /// 协议 ID（对齐 C++ `NN_ProtoCmd_*`）
    pub proto_id: u32,
    /// 序列号（和 Response 配对，供 client 端请求-响应匹配）
    pub serial_no: u32,
    /// 请求体 proto 格式（Protobuf / JSON）
    pub proto_fmt_type: ProtoFmtType,
    /// 请求 body（已解密后的明文）
    pub body: Bytes,
    /// v1.4.38 Phase 4: 订单幂等 key（由 REST `Idempotency-Key` header / gRPC
    /// metadata / WS envelope / MCP tool args 填入）。None 表示客户端未传，
    /// handler 走无幂等直通 path（backward-compat）。
    pub idempotency_key: Option<String>,
    /// v1.4.106 codex 0920 F1 (P1): caller key id 副本 (per-call snapshot,
    /// 由 surface adapter 层从 KeyRecord 读取后填入).
    ///
    /// **目标**: idempotency cache key namespace 必须含 caller key id, 否则
    /// 不同 caller 用同 Idempotency-Key 会跨 caller 命中老 response —— 严重
    /// 跨账户数据泄漏 + 重复下单 silent fail.
    ///
    /// `None` = 无 caller 标识 (legacy TCP / 未 auth) → namespace 用 `<no_key>`
    /// 占位符. `Some("alice")` = WS / MCP / REST 已 auth 的 caller —— namespace
    /// 用 `<caller_key_id="alice">`, 不与其他 caller 串.
    pub caller_key_id: Option<String>,
    /// v1.4.105 D2 contract-hardening 补丁: caller key 的 `allowed_acc_ids` 硬限额
    /// 副本 (per-call snapshot, 在 surface adapter 层从 KeyRecord 读取后填入).
    ///
    /// **目标**: 让 dispatch-time handlers (e.g. `SubAccPushHandler` 注册 acc_id
    /// 到 SubscriptionManager) 也能 enforce per-acc whitelist — 即使上游 pipeline
    /// body-aware step 已 enforce, 让 handler 自己 defense-in-depth 防 future
    /// regression (新 surface 加进来漏调 pipeline body-aware).
    ///
    /// `None` / `Some(empty)` = caller 无 acc_id 限制 (legacy mode 或 unrestricted
    /// key) → handler 不 filter; `Some(non-empty set)` → handler 应 reject 不在
    /// set 中的 acc_id. Deny-all 使用 sentinel `{0}`，不使用空集合。
    pub caller_allowed_acc_ids: Option<std::sync::Arc<std::collections::HashSet<u64>>>,
}

impl IncomingRequest {
    /// codex 0522 F4 v1.4.106: cross-surface 单测 hook. 构 IncomingRequest
    /// 并填 caller scope (`caller_key_id` + `caller_allowed_acc_ids`) — 让
    /// REST / gRPC / WS / MCP 等 surface 的 adapter 都用同一构造路径, 防
    /// "某个 surface 漏填字段" silent regression.
    ///
    /// 之前 4 surface 各写一份 struct literal, 加新字段需逐个改, 漏一个就
    /// 出现 silent None — 与坑 #54 schema-only fix 同模式 (实装符号 vs 真
    /// 行为差距). 本 helper 是 single point, 加新字段 schema 自动 propagate.
    ///
    /// **注意**: 本 helper 不 take ownership of body — caller 已 own bytes.
    /// idempotency_key / caller_key_id 接 String 而非 &str 让 caller 决定
    /// 是 clone 还是 move.
    pub fn builder(
        conn_id: u64,
        proto_id: u32,
        serial_no: u32,
        proto_fmt_type: ProtoFmtType,
        body: Bytes,
    ) -> IncomingRequestBuilder {
        IncomingRequestBuilder {
            request: Self {
                conn_id,
                proto_id,
                serial_no,
                proto_fmt_type,
                body,
                idempotency_key: None,
                caller_allowed_acc_ids: None,
                caller_key_id: None,
            },
        }
    }
}

/// Thin builder for `IncomingRequest`.
///
/// The base request shape is the wire envelope; idempotency and caller scope are
/// optional per-surface decorations. Keeping those defaults here avoids every
/// REST / gRPC / raw WS / MCP adapter spelling out `None` independently.
#[derive(Debug)]
pub struct IncomingRequestBuilder {
    request: IncomingRequest,
}

impl IncomingRequestBuilder {
    pub fn with_idempotency_key(mut self, idempotency_key: Option<String>) -> Self {
        self.request.idempotency_key = idempotency_key;
        self
    }

    pub fn with_caller_scope(
        mut self,
        caller_allowed_acc_ids: Option<std::sync::Arc<HashSet<u64>>>,
        caller_key_id: Option<String>,
    ) -> Self {
        self.request.caller_allowed_acc_ids = caller_allowed_acc_ids;
        self.request.caller_key_id = caller_key_id;
        self
    }

    pub fn build(self) -> IncomingRequest {
        self.request
    }
}

impl From<IncomingRequestBuilder> for IncomingRequest {
    fn from(builder: IncomingRequestBuilder) -> Self {
        builder.build()
    }
}

impl ClientConn {
    /// 生成随机连接 ID（与 C++ 的 GetRand_MilliTimeAndU22 对应）
    pub fn generate_conn_id() -> u64 {
        use std::time::{SystemTime, UNIX_EPOCH};
        let millis = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        let rand_part: u32 = rand::random();
        (millis << 22) | (rand_part as u64 & 0x3FFFFF)
    }

    /// 生成随机 AES key（16 字节 hex 字符串的 ASCII 字节）
    pub fn generate_aes_key() -> [u8; 16] {
        let rand_val: u64 = rand::random();
        let hex = format!("{rand_val:016X}");
        let mut key = [0u8; 16];
        key.copy_from_slice(hex.as_bytes());
        key
    }

    /// 创建发送帧，自动处理 AES 加密
    ///
    /// 当 aes_encrypt_enabled 为 true 时：
    /// - SHA1 基于明文计算（FutuFrame::new 自动处理）
    /// - body 使用 AES-128 ECB 加密
    /// - header.body_len 更新为密文长度
    ///
    /// 对应 C++ APIServerCS_Conn::OnSendPacketData 的加密逻辑
    pub fn make_frame(&self, proto_id: u32, serial_no: u32, body: Bytes) -> FutuFrame {
        if self.aes_encrypt_enabled {
            // 先用明文构建 frame（SHA1 基于明文计算）
            let mut frame = FutuFrame::new(proto_id, serial_no, body);
            // 再加密 body
            let encrypted = encrypt::aes_ecb_encrypt(&self.aes_key, &frame.body);
            frame.header.body_len = encrypted.len() as u32;
            frame.body = Bytes::from(encrypted);
            frame
        } else {
            FutuFrame::new(proto_id, serial_no, body)
        }
    }

    /// 解密请求 body（如果启用了 AES 加密）
    ///
    /// 对应 C++ APIServerCS_Conn::OnRecvPacket 的解密逻辑
    pub fn decrypt_body(&self, body: &[u8]) -> Result<Vec<u8>, FutuError> {
        if self.aes_encrypt_enabled {
            encrypt::aes_ecb_decrypt(&self.aes_key, body).map_err(|e| {
                tracing::warn!(conn_id = self.conn_id, error = %e, "AES decrypt body failed");
                e
            })
        } else {
            Ok(body.to_vec())
        }
    }

    /// 处理 InitConnect 请求，返回 InitConnect 响应 body
    ///
    /// 当配置了 RSA 私钥时：
    /// - C2S 请求 body 使用 RSA 公钥加密（需要用私钥解密）
    /// - S2C 响应 body 使用 RSA 公钥加密（客户端用私钥解密）
    ///
    /// 对应 C++ APIServer::OnRecvInitConnect
    pub fn handle_init_connect(
        &mut self,
        body: &[u8],
        server_ver: i32,
        login_user_id: u64,
        keepalive_interval: i32,
        rsa_private_key: Option<&str>,
    ) -> Result<Vec<u8>, FutuError> {
        // 1. 解密 C2S（如果配置了 RSA）
        let decrypted_body;
        let req_body = if let Some(rsa_key) = rsa_private_key {
            decrypted_body =
                futu_net::encrypt::rsa_private_decrypt_blocks(rsa_key, body).map_err(|e| {
                    tracing::warn!(error = %e, "RSA decrypt InitConnect C2S failed");
                    e
                })?;
            tracing::debug!(
                encrypted_len = body.len(),
                decrypted_len = decrypted_body.len(),
                "RSA decrypted InitConnect C2S"
            );
            &decrypted_body[..]
        } else {
            body
        };

        let req: futu_proto::init_connect::Request =
            prost::Message::decode(req_body).map_err(FutuError::Proto)?;

        self.state = ConnState::Initialized;
        self.recv_notify = req.c2s.recv_notify.unwrap_or(false);

        // 当配置了 RSA 时，后续所有帧使用 AES 加解密
        if rsa_private_key.is_some() {
            self.aes_encrypt_enabled = true;
            tracing::debug!(conn_id = self.conn_id, "AES body encryption enabled");
        }

        let aes_key_str = std::str::from_utf8(&self.aes_key)
            .map_err(|e| FutuError::Codec(format!("invalid InitConnect conn_aes_key: {e}")))?
            .to_string();

        let resp = futu_proto::init_connect::Response {
            ret_type: 0,
            ret_msg: None,
            err_code: None,
            s2c: Some(futu_proto::init_connect::S2c {
                server_ver,
                login_user_id,
                conn_id: self.conn_id,
                conn_aes_key: aes_key_str,
                keep_alive_interval: keepalive_interval,
                aes_cb_civ: None,
                user_attribution: None,
            }),
        };

        let resp_body = prost::Message::encode_to_vec(&resp);

        // 2. 加密 S2C（如果配置了 RSA）
        if let Some(rsa_key) = rsa_private_key {
            let encrypted = futu_net::encrypt::rsa_public_encrypt_blocks(rsa_key, &resp_body)
                .map_err(|e| {
                    tracing::warn!(error = %e, "RSA encrypt InitConnect S2C failed");
                    e
                })?;
            tracing::debug!(
                plaintext_len = resp_body.len(),
                encrypted_len = encrypted.len(),
                "RSA encrypted InitConnect S2C"
            );
            Ok(encrypted)
        } else {
            Ok(resp_body)
        }
    }

    /// 处理 KeepAlive 请求
    pub fn handle_keepalive(&self, body: &[u8]) -> Result<Vec<u8>, FutuError> {
        let _req: futu_proto::keep_alive::Request =
            prost::Message::decode(body).map_err(FutuError::Proto)?;

        self.keepalive_count.fetch_add(1, Ordering::Relaxed);

        let resp = futu_proto::keep_alive::Response {
            ret_type: 0,
            ret_msg: None,
            err_code: None,
            s2c: Some(futu_proto::keep_alive::S2c {
                time: chrono::Utc::now().timestamp(),
            }),
        };

        Ok(prost::Message::encode_to_vec(&resp))
    }
}

/// 连接断开通知
pub struct DisconnectNotify {
    /// 被断开的连接 ID（订阅 / push / auth 状态清理用）
    pub conn_id: u64,
}

/// 运行单个连接的收发循环
///
/// 返回 (接收请求的 channel, 连接信息)
pub async fn run_connection(
    stream: TcpStream,
    conn_id: u64,
    _aes_key: [u8; 16],
    req_tx: mpsc::UnboundedSender<IncomingRequest>,
    disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
) -> mpsc::Sender<FutuFrame> {
    let (frame_tx, mut frame_rx) = mpsc::channel::<FutuFrame>(256);

    let framed = Framed::new(stream, FutuCodec);
    let (mut sink, mut stream) = framed.split();

    // 发送任务
    tokio::spawn(async move {
        while let Some(frame) = frame_rx.recv().await {
            if let Err(e) = sink.send(frame).await {
                tracing::warn!(conn_id = conn_id, error = %e, "send failed");
                break;
            }
        }
    });

    // 接收任务
    tokio::spawn(async move {
        while let Some(result) = stream.next().await {
            match result {
                Ok(frame) => {
                    // TCP wire currently has no idempotency or key scope fields.
                    // The builder keeps legacy defaults explicit in one place.
                    let req = IncomingRequest::builder(
                        conn_id,
                        frame.header.proto_id,
                        frame.header.serial_no,
                        frame.header.proto_fmt_type,
                        frame.body,
                    )
                    .build();
                    if req_tx.send(req).is_err() {
                        break;
                    }
                }
                Err(e) => {
                    tracing::warn!(conn_id = conn_id, error = %e, "recv error");
                    break;
                }
            }
        }
        tracing::info!(conn_id = conn_id, "connection closed");
        // 通知 listener 清理连接
        let _ = disconnect_tx.send(DisconnectNotify { conn_id });
    });

    frame_tx
}

#[cfg(test)]
mod tests;