futu_server/

listener.rs

// TCP 监听器：接受客户端连接，管理连接池

use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};
use std::time::Instant;

use dashmap::DashMap;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, watch};

use futu_codec::header::ProtoFmtType;
use futu_core::proto_id;

use crate::conn::{ClientConn, ConnState, DisconnectNotify, IncomingRequest};
use crate::metrics::GatewayMetrics;
use crate::router::RequestRouter;

/// 服务端最大连接数
pub const MAX_CONNECTIONS: usize = 128;

/// Inbound request queue capacity shared by raw TCP and WebSocket listeners.
///
/// A bounded queue turns slow backend dispatch into socket-level backpressure
/// instead of letting client frames accumulate without a memory ceiling.
pub(crate) const REQUEST_QUEUE_CAPACITY: usize = 4096;

/// 服务端配置
#[derive(Debug, Clone)]
pub struct ServerConfig {
    /// TCP 监听地址（如 `127.0.0.1:11111`）
    pub listen_addr: String,
    /// 服务端版本号，InitConnect 响应下发给客户端
    pub server_ver: i32,
    /// 服务端登录 user_id，InitConnect 响应下发给客户端
    pub login_user_id: u64,
    /// KeepAlive 心跳间隔（秒），InitConnect 响应下发给客户端
    pub keepalive_interval: i32,
    /// RSA 私钥 PEM 内容（可选，启用后 InitConnect 使用 RSA 加解密）
    pub rsa_private_key: Option<String>,
}

#[must_use]
pub(crate) fn default_server_time_offset_secs() -> Arc<AtomicI64> {
    Arc::new(AtomicI64::new(0))
}

#[must_use]
pub(crate) fn server_now_ts_at(server_time_offset_secs: &AtomicI64, local_now_ts: i64) -> i64 {
    local_now_ts.saturating_add(server_time_offset_secs.load(Ordering::Relaxed))
}

#[must_use]
pub(crate) fn server_now_ts(server_time_offset_secs: &AtomicI64) -> i64 {
    server_now_ts_at(server_time_offset_secs, chrono::Utc::now().timestamp())
}

pub(crate) fn set_nodelay_with_log(
    stream: &tokio::net::TcpStream,
    peer_addr: std::net::SocketAddr,
    surface: &'static str,
) {
    if let Err(error) = stream.set_nodelay(true) {
        tracing::debug!(
            peer = %peer_addr,
            surface,
            error = %error,
            "tcp nodelay setup failed"
        );
    }
}

pub(crate) async fn shutdown_requested(shutdown_rx: &mut watch::Receiver<bool>) {
    loop {
        if *shutdown_rx.borrow() {
            return;
        }
        if shutdown_rx.changed().await.is_err() {
            return;
        }
    }
}

/// API 服务端
pub struct ApiServer {
    config: ServerConfig,
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
    metrics: Arc<GatewayMetrics>,
    server_time_offset_secs: Arc<AtomicI64>,
}

impl ApiServer {
    /// 创建新的服务端实例。不自动启动，需调用 [`ApiServer::run`] 进入接收循环。
    pub fn new(config: ServerConfig) -> Self {
        Self {
            config,
            connections: Arc::new(DashMap::new()),
            router: Arc::new(RequestRouter::new()),
            subscriptions: None,
            metrics: Arc::new(GatewayMetrics::new()),
            server_time_offset_secs: default_server_time_offset_secs(),
        }
    }

    /// 设置订阅管理器，用于连接断开时自动清理订阅关系
    pub fn set_subscriptions(&mut self, subs: Arc<crate::subscription::SubscriptionManager>) {
        self.subscriptions = Some(subs);
    }

    /// 获取路由器引用（用于注册业务处理器）
    pub fn router(&self) -> &Arc<RequestRouter> {
        &self.router
    }

    /// 获取连接池引用（用于推送分发）
    pub fn connections(&self) -> &Arc<DashMap<u64, ClientConn>> {
        &self.connections
    }

    /// 设置外部监控指标（共享同一个 Arc，让 bridge 和 server 使用同一份计数器）
    pub fn set_metrics(&mut self, metrics: Arc<GatewayMetrics>) {
        self.metrics = metrics;
    }

    /// Inject backend server-time offset for SDK-facing API protocol fields.
    pub fn set_server_time_offset_secs(&mut self, offset: Arc<AtomicI64>) {
        self.server_time_offset_secs = offset;
    }

    /// 获取监控指标引用
    pub fn metrics(&self) -> &Arc<GatewayMetrics> {
        &self.metrics
    }

    /// 启动服务端监听
    pub async fn run(&self) -> anyhow::Result<()> {
        let (_shutdown_tx, shutdown_rx) = watch::channel(false);
        self.run_until_shutdown(shutdown_rx).await
    }

    /// 启动服务端监听，并在 shutdown 信号到来时停止接受新连接。
    pub async fn run_until_shutdown(
        &self,
        mut shutdown_rx: watch::Receiver<bool>,
    ) -> anyhow::Result<()> {
        let listener = TcpListener::bind(&self.config.listen_addr).await?;
        tracing::info!(addr = %self.config.listen_addr, "API server listening");

        let (req_tx, req_rx) = mpsc::channel::<IncomingRequest>(REQUEST_QUEUE_CAPACITY);
        // The disconnect cleanup signal is intentionally unbounded: each item
        // is a tiny `conn_id`, and blocking this path behind request-queue
        // backpressure would risk leaking connection/subscription state.
        let (disconnect_tx, mut disconnect_rx) = mpsc::unbounded_channel::<DisconnectNotify>();

        // 启动请求处理任务
        let connections = Arc::clone(&self.connections);
        let router = Arc::clone(&self.router);
        let config = self.config.clone();
        let metrics = Arc::clone(&self.metrics);
        let server_time_offset_secs = Arc::clone(&self.server_time_offset_secs);
        tokio::spawn(async move {
            process_requests(
                req_rx,
                connections,
                router,
                config,
                metrics,
                server_time_offset_secs,
            )
            .await;
        });

        // 启动连接清理任务（TCP 断开通知）
        let cleanup_connections = Arc::clone(&self.connections);
        let cleanup_subs = self.subscriptions.clone();
        let cleanup_metrics = Arc::clone(&self.metrics);
        tokio::spawn(async move {
            while let Some(notify) = disconnect_rx.recv().await {
                let removed = cleanup_connections.remove(&notify.conn_id);
                if removed.is_some() {
                    cleanup_metrics
                        .total_disconnections
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    // 清理该连接的所有订阅关系
                    if let Some(ref subs) = cleanup_subs {
                        subs.on_disconnect(notify.conn_id);
                    }
                    tracing::info!(
                        conn_id = notify.conn_id,
                        remaining = cleanup_connections.len(),
                        "connection removed from pool"
                    );
                }
            }
        });

        // 启动 KeepAlive 超时检测任务（对应 C++ OnTimeTicker，每 66 秒无活动断连）
        let ka_connections = Arc::clone(&self.connections);
        let ka_subs = self.subscriptions.clone();
        let ka_metrics = Arc::clone(&self.metrics);
        let mut ka_shutdown_rx = shutdown_rx.clone();
        tokio::spawn(async move {
            const CHECK_INTERVAL_SECS: u64 = 15;
            const TIMEOUT_SECS: u64 = 66;
            let mut interval =
                tokio::time::interval(std::time::Duration::from_secs(CHECK_INTERVAL_SECS));
            interval.tick().await; // 跳过首次立即触发
            loop {
                tokio::select! {
                    _ = shutdown_requested(&mut ka_shutdown_rx) => {
                        tracing::info!("API server keepalive task stopped by shutdown signal");
                        break;
                    }
                    _ = interval.tick() => {}
                }
                let now = Instant::now();
                let mut timed_out = Vec::new();
                for entry in ka_connections.iter() {
                    let conn = entry.value();
                    if now.duration_since(conn.last_keepalive).as_secs() >= TIMEOUT_SECS {
                        timed_out.push(conn.conn_id);
                    }
                }
                for conn_id in timed_out {
                    if ka_connections.remove(&conn_id).is_some() {
                        ka_metrics
                            .keepalive_timeouts
                            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                        ka_metrics
                            .total_disconnections
                            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                        if let Some(ref subs) = ka_subs {
                            subs.on_disconnect(conn_id);
                        }
                        tracing::info!(
                            conn_id = conn_id,
                            remaining = ka_connections.len(),
                            "keepalive timeout, connection removed"
                        );
                    }
                }
            }
        });

        // 接受连接循环
        let connections = Arc::clone(&self.connections);
        let accept_metrics = Arc::clone(&self.metrics);
        loop {
            let (stream, peer_addr) = tokio::select! {
                _ = shutdown_requested(&mut shutdown_rx) => {
                    tracing::info!("API server accept loop stopped by shutdown signal");
                    break;
                }
                accepted = listener.accept() => accepted?,
            };

            if connections.len() >= MAX_CONNECTIONS {
                accept_metrics
                    .rejected_connections
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                tracing::warn!(
                    peer = %peer_addr,
                    "max connections reached ({}), rejecting",
                    MAX_CONNECTIONS
                );
                drop(stream);
                continue;
            }

            accept_metrics
                .total_connections
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

            let conn_id = crate::conn::ClientConn::generate_conn_id();
            let aes_key = crate::conn::ClientConn::generate_aes_key();
            set_nodelay_with_log(&stream, peer_addr, "tcp");

            tracing::info!(
                conn_id = conn_id,
                peer = %peer_addr,
                total = connections.len() + 1,
                "client connected"
            );

            let tx = crate::conn::run_connection(
                stream,
                conn_id,
                aes_key,
                req_tx.clone(),
                disconnect_tx.clone(),
                shutdown_rx.clone(),
            )
            .await;

            let conn = ClientConn {
                conn_id,
                state: ConnState::Connected,
                aes_key,
                aes_encrypt_enabled: false,
                proto_fmt_type: ProtoFmtType::Protobuf,
                last_keepalive: Instant::now(),
                recv_notify: false,
                ai_type: 0,
                keepalive_count: std::sync::atomic::AtomicU32::new(0),
                tx,
                // 原 TCP listener 不做 per-message scope 校验（保持兼容）：
                // key_id=None / scopes=空集 被 ws_listener 的 gate 解释为"legacy 全放行"
                key_id: None,
                scopes: std::collections::HashSet::new(),
                // v1.4.105 D3 (Phase 4) T-B2: TCP listener 同样 legacy 模式 →
                // allowed_markets None = 无限制 (push_trd_acc Layer 3 不 trigger).
                allowed_markets: None,
                // codex round 1 F4 (P2) v1.4.105: 同 legacy 模式, 无 acc_id 限制.
                allowed_acc_ids: None,
            };

            connections.insert(conn_id, conn);
        }

        Ok(())
    }

    /// 向指定连接发送响应（自动处理 AES 加密）
    pub async fn send_response(
        connections: &DashMap<u64, ClientConn>,
        conn_id: u64,
        proto_id: u32,
        serial_no: u32,
        body: Vec<u8>,
    ) {
        if let Some(conn) = connections.get(&conn_id) {
            let frame = conn.make_frame(proto_id, serial_no, bytes::Bytes::from(body));
            if conn.tx.send(frame).await.is_err() {
                tracing::warn!(
                    conn_id = conn_id,
                    "failed to send response, connection closed"
                );
            }
        }
    }
}

/// 处理所有连接的请求
async fn process_requests(
    mut req_rx: mpsc::Receiver<IncomingRequest>,
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    config: ServerConfig,
    metrics: Arc<GatewayMetrics>,
    server_time_offset_secs: Arc<AtomicI64>,
) {
    while let Some(mut req) = req_rx.recv().await {
        let conn_id = req.conn_id;
        let proto_id_val = req.proto_id;
        let serial_no = req.serial_no;
        let req_start = Instant::now();

        metrics
            .total_requests
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        // 更新 last_keepalive（任何包都算活跃，对应 C++ m_nKeepAlive_Count_Curt++）
        if let Some(mut conn) = connections.get_mut(&conn_id) {
            conn.last_keepalive = Instant::now();
        }

        // v1.4.106 codex 0532 F3 (P2): daemon-internal proto_id (高位
        // 0x8000_0000 bit) 绝不应从 raw TCP 公开 surface 进入 — 仅 REST
        // handler 内部合成给 router. 显式 reject + log, 防探测 daemon
        // 内部 routing.
        if futu_auth::is_internal_proto_id(proto_id_val) {
            metrics
                .total_request_errors
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            tracing::warn!(
                conn_id,
                proto_id = proto_id_val,
                "rejecting daemon-internal proto_id at raw TCP public surface (audit 0532 F3)"
            );
            continue;
        }

        // 非 InitConnect 请求需要 AES 解密（InitConnect 自身处理 RSA 解密）
        if proto_id_val != proto_id::INIT_CONNECT
            && let Some(conn) = connections.get(&conn_id)
            && conn.aes_encrypt_enabled
        {
            match conn.decrypt_body(&req.body) {
                Ok(decrypted) => {
                    req.body = bytes::Bytes::from(decrypted);
                }
                Err(e) => {
                    metrics
                        .total_request_errors
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    tracing::warn!(
                        conn_id = conn_id,
                        proto_id = proto_id_val,
                        error = %e,
                        "AES decrypt request failed, dropping"
                    );
                    continue;
                }
            }
        }

        // InitConnect 和 KeepAlive 内部处理
        let response_body = futu_backend::delay_stats::with_api_request(
            conn_id,
            serial_no,
            proto_id_val,
            || async {
                match proto_id_val {
                    proto_id::INIT_CONNECT => match connections.get_mut(&conn_id) {
                        Some(mut conn) => match conn.handle_init_connect(
                            &req.body,
                            config.server_ver,
                            config.login_user_id,
                            config.keepalive_interval,
                            config.rsa_private_key.as_deref(),
                        ) {
                            Ok(body) => Some(body),
                            Err(error) => {
                                metrics
                                    .total_request_errors
                                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                                tracing::warn!(
                                    conn_id,
                                    proto_id = proto_id_val,
                                    error = %error,
                                    "InitConnect handling failed"
                                );
                                None
                            }
                        },
                        None => {
                            metrics
                                .total_request_errors
                                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                            tracing::warn!(
                                conn_id,
                                proto_id = proto_id_val,
                                "InitConnect request received for missing connection"
                            );
                            None
                        }
                    },
                    proto_id::KEEP_ALIVE => match connections.get(&conn_id) {
                        Some(conn) => match conn
                            .handle_keepalive_at(&req.body, server_now_ts(&server_time_offset_secs))
                        {
                            Ok(body) => Some(body),
                            Err(error) => {
                                metrics
                                    .total_request_errors
                                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                                tracing::warn!(
                                    conn_id,
                                    proto_id = proto_id_val,
                                    error = %error,
                                    "KeepAlive handling failed"
                                );
                                None
                            }
                        },
                        None => {
                            metrics
                                .total_request_errors
                                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                            tracing::warn!(
                                conn_id,
                                proto_id = proto_id_val,
                                "KeepAlive request received for missing connection"
                            );
                            None
                        }
                    },
                    _ => {
                        // 委托给路由器
                        router.dispatch(conn_id, &req).await
                    }
                }
            },
        )
        .await;

        // 记录延迟
        metrics.record_latency_ns(req_start.elapsed().as_nanos() as u64);

        // 发送响应
        if let Some(body) = response_body {
            metrics
                .total_response_bytes
                .fetch_add(body.len() as u64, std::sync::atomic::Ordering::Relaxed);
            ApiServer::send_response(&connections, conn_id, proto_id_val, serial_no, body).await;
        } else if proto_id_val != proto_id::INIT_CONNECT && proto_id_val != proto_id::KEEP_ALIVE {
            metrics
                .total_request_errors
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
    }
}

#[cfg(test)]
mod tests;