futu_server/

listener.rs

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

use std::sync::Arc;
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 futu_core::server_time::ServerTimeAnchorStore;

use crate::conn::{ClientConn, ConnState, DisconnectNotify, IncomingRequest};
use crate::listener_status::{
    ListenerBindEventSender, ListenerSurface, notify_listener_failed, notify_listener_opened,
};
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;

/// Per-client share of the existing listener-wide inbound request budget.
///
/// A dedicated FIFO worker prevents a slow backend dispatch on one connection
/// from blocking InitConnect or business traffic on another connection.
pub(crate) const PER_CONNECTION_REQUEST_QUEUE_CAPACITY: usize =
    REQUEST_QUEUE_CAPACITY / MAX_CONNECTIONS;
const _: () = assert!(REQUEST_QUEUE_CAPACITY.is_multiple_of(MAX_CONNECTIONS));
const _: () = assert!(PER_CONNECTION_REQUEST_QUEUE_CAPACITY > 0);

/// 服务端配置
#[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_store() -> Arc<ServerTimeAnchorStore> {
    Arc::new(ServerTimeAnchorStore::new())
}

#[must_use]
pub(crate) fn server_now_ts_at(
    server_time_store: &ServerTimeAnchorStore,
    local_now_unix_micros: i64,
    monotonic_now: Instant,
) -> i64 {
    server_time_store
        .now_unix_micros_at(local_now_unix_micros, monotonic_now)
        .div_euclid(1_000_000)
}

#[must_use]
pub(crate) fn server_now_ts(server_time_store: &ServerTimeAnchorStore) -> i64 {
    server_now_ts_at(
        server_time_store,
        chrono::Utc::now().timestamp_micros(),
        Instant::now(),
    )
}

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_store: Arc<ServerTimeAnchorStore>,
}

impl ApiServer {
    /// 创建新的服务端实例。不自动启动，需调用 [`ApiServer::run`] 进入接收循环。
    pub fn new(config: ServerConfig) -> Self {
        let startup_readiness =
            crate::identity::StartupReadiness::ready(config.login_user_id, None);
        Self {
            config,
            connections: Arc::new(DashMap::new()),
            router: Arc::new(RequestRouter::with_startup_readiness(startup_readiness)),
            subscriptions: None,
            metrics: Arc::new(GatewayMetrics::new()),
            server_time_store: default_server_time_store(),
        }
    }

    /// 设置订阅管理器，用于连接断开时自动清理订阅关系
    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 the shared backend server-time anchor for SDK-facing fields.
    pub fn set_server_time_store(&mut self, store: Arc<ServerTimeAnchorStore>) {
        self.server_time_store = store;
    }

    /// 获取监控指标引用
    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,
        shutdown_rx: watch::Receiver<bool>,
    ) -> anyhow::Result<()> {
        self.run_until_shutdown_with_listener_events(shutdown_rx, None)
            .await
    }

    /// Run until shutdown and report the exact socket bind result to startup.
    pub async fn run_until_shutdown_with_listener_events(
        &self,
        mut shutdown_rx: watch::Receiver<bool>,
        listener_events: Option<ListenerBindEventSender>,
    ) -> anyhow::Result<()> {
        let listener = TcpListener::bind(&self.config.listen_addr)
            .await
            .map_err(|error| {
                notify_listener_failed(&listener_events, ListenerSurface::Ftapi);
                anyhow::Error::new(crate::bind_hint::io_bind_error(
                    "FTAPI TCP",
                    "--port",
                    &self.config.listen_addr,
                    error,
                ))
            })?;
        tracing::info!(addr = %self.config.listen_addr, "API server listening");

        // 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>();

        // 启动连接清理任务（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);
        let _serving =
            notify_listener_opened(&listener_events, ListenerSurface::Ftapi, &shutdown_rx).await?;
        drop(listener_events);

        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 session_generation = crate::conn::ClientConn::generate_session_generation();
            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 (req_tx, req_rx) =
                mpsc::channel::<IncomingRequest>(PER_CONNECTION_REQUEST_QUEUE_CAPACITY);
            let (tx, close_control, io_start) = crate::conn::prepare_connection_with_close_control(
                stream,
                conn_id,
                aes_key,
                req_tx,
                disconnect_tx.clone(),
                shutdown_rx.clone(),
            )
            .await;

            let conn = ClientConn {
                conn_id,
                session_generation,
                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);
            if let Some(ref subscriptions) = self.subscriptions {
                subscriptions.on_connect(conn_id, session_generation);
                subscriptions.register_client_close_control(conn_id, close_control);
            }
            tokio::spawn(process_requests(
                req_rx,
                Arc::clone(&connections),
                Arc::clone(&self.router),
                self.subscriptions.clone(),
                self.config.clone(),
                Arc::clone(&self.metrics),
                Arc::clone(&self.server_time_store),
                peer_addr.ip().is_loopback(),
            ));
            io_start.start();
        }

        Ok(())
    }

    /// 向指定连接发送响应（自动处理 AES 加密）
    pub async fn send_response(
        connections: &DashMap<u64, ClientConn>,
        conn_id: u64,
        proto_id: u32,
        serial_no: u32,
        body: Vec<u8>,
    ) -> bool {
        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"
                );
                return false;
            }
            return true;
        }
        false
    }
}

pub(crate) fn reconcile_post_dispatch_connection_state(
    connections: &DashMap<u64, ClientConn>,
    subscriptions: Option<&crate::subscription::SubscriptionManager>,
    conn_id: u64,
) {
    if connections.contains_key(&conn_id) {
        return;
    }

    if let Some(subscriptions) = subscriptions {
        let became_empty = subscriptions.on_disconnect(conn_id);
        tracing::debug!(
            conn_id,
            became_empty_count = became_empty.len(),
            "reconciled subscription state committed after connection disconnect"
        );
    }
}

/// 处理所有连接的请求
async fn process_requests(
    mut req_rx: mpsc::Receiver<IncomingRequest>,
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
    config: ServerConfig,
    metrics: Arc<GatewayMetrics>,
    server_time_store: Arc<ServerTimeAnchorStore>,
    peer_is_loopback: bool,
) {
    let mut pending_init_connect: Option<tokio::task::JoinHandle<()>> = None;
    while let Some(mut req) = req_rx.recv().await {
        req.caller_is_loopback = peer_is_loopback;
        req.caller_legacy_local_mode = true;
        let conn_id = req.conn_id;
        let Some(session_generation) = connections
            .get(&conn_id)
            .map(|conn| conn.session_generation)
        else {
            continue;
        };
        req.session_generation = session_generation;
        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;
        }

        if proto_id_val != proto_id::INIT_CONNECT
            && connections
                .get(&conn_id)
                .is_some_and(|conn| conn.state == crate::conn::ConnState::Connected)
            && !crate::identity::StartupReadiness::is_prelogin_proto(proto_id_val)
        {
            if router.startup_readiness().snapshot().state != crate::identity::StartupState::Ready {
                if let Some(body) = router.dispatch(conn_id, &req).await {
                    metrics
                        .total_response_bytes
                        .fetch_add(body.len() as u64, std::sync::atomic::Ordering::Relaxed);
                    if ApiServer::send_response(
                        connections.as_ref(),
                        conn_id,
                        proto_id_val,
                        serial_no,
                        body,
                    )
                    .await
                    {
                        req.mark_response_committed();
                    }
                }
                continue;
            }
            let Some(pending) = pending_init_connect.take() else {
                tracing::warn!(
                    conn_id,
                    proto_id = proto_id_val,
                    "dropping request from Ready connection with no completed InitConnect"
                );
                continue;
            };
            if let Err(error) = pending.await {
                tracing::warn!(conn_id, error = %error, "pending InitConnect task failed");
                continue;
            }
            if connections
                .get(&conn_id)
                .is_none_or(|conn| conn.state == crate::conn::ConnState::Connected)
            {
                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_core::delay_stats::with_api_request(conn_id, serial_no, proto_id_val, || async {
                match proto_id_val {
                    proto_id::INIT_CONNECT => match crate::conn::start_init_connect_for_startup(
                        Arc::clone(&connections),
                        router.startup_readiness().clone(),
                        conn_id,
                        &req.body,
                        serial_no,
                        config.server_ver,
                        config.keepalive_interval,
                        config.rsa_private_key.clone(),
                    ) {
                        Ok(crate::conn::InitConnectStart::Immediate(prepared)) => {
                            if let Err(error) = crate::conn::complete_prepared_init_connect(
                                connections.as_ref(),
                                conn_id,
                                serial_no,
                                prepared,
                            )
                            .await
                            {
                                metrics
                                    .total_request_errors
                                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                                tracing::warn!(conn_id, error = %error, "InitConnect response failed");
                            }
                            None
                        }
                        Ok(crate::conn::InitConnectStart::Deferred(task)) => {
                            if let Some(previous) = pending_init_connect.replace(task) {
                                previous.abort();
                            }
                            None
                        }
                        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
                        }
                    },
                    proto_id::KEEP_ALIVE => match connections.get(&conn_id) {
                        Some(conn) => match conn
                            .handle_keepalive_at(&req.body, server_now_ts(&server_time_store))
                        {
                            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);

        reconcile_post_dispatch_connection_state(&connections, subscriptions.as_deref(), conn_id);

        // 发送响应
        if let Some(body) = response_body {
            metrics
                .total_response_bytes
                .fetch_add(body.len() as u64, std::sync::atomic::Ordering::Relaxed);
            if ApiServer::send_response(&connections, conn_id, proto_id_val, serial_no, body).await
            {
                req.mark_response_committed();
            }
        } 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);
        }
    }
    if let Some(pending) = pending_init_connect {
        pending.abort();
    }
}

#[cfg(test)]
mod tests;