futu_server/

ws_listener.rs

// WebSocket 监听器：接受 WebSocket 连接，复用 TCP 的请求路由和连接池
//
// 每个 WebSocket 二进制消息 = 一个完整的 FutuAPI 帧（44 字节帧头 + body）。
// 与 TCP 共享同一个 connections DashMap、RequestRouter、SubscriptionManager。
//
// ## v1.0 鉴权
//
// 握手阶段（accept_hdr_async）校验 HTTP `Authorization: Bearer <token>` 或
// `?token=<plaintext>` query —— 通过 `KeyStore::verify` 得到 `KeyRecord`，
// 把 scope 集合和 key_id 存到 `ClientConn`。每条消息进 `ws_process_requests`
// 时按 `futu_auth::scope_for_proto_id(proto_id)` 查所需 scope，不匹配 → 不 dispatch、
// 记 audit reject。`trade:real` 额外跑 `check_and_commit` 过一道 rate + hours
// 全局闸门。未注入 KeyStore（TCP listener 或 legacy 模式）→ scopes 空集被
// 解释为"全放行"，保持向后兼容。
//
// ## v1.4.104 阶段 3
//
// 把 inline scope gate / rate gate / body-aware acc_id check 替换为
// `futu_auth_pipeline::authenticate_request` 单一调用. 流程:
//   1. AES decrypt (handshake protocol INIT_CONNECT 跳过)
//   2. (非 INIT_CONNECT + 非 1xxx 系统协议) → 调 pipeline
//      Credential::PreVerified(rec_from_get_by_id) 拿 SIGHUP-aware 最新 rec.
//      Reject → drop. Allow → 取 allowed_acc_ids 给 response filter.
//   3. dispatch (router / handle_init_connect / handle_keepalive)
//   4. response filter (proto 2001 TRD_GET_ACC_LIST 等)
//
// 删除: 旧 `ws_body_aware_check` (功能进 pipeline body_aware::build_check_ctxs).
// 删除: 旧 inline scope gate + rate gate + audit allow/reject. 全 pipeline 一处.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Instant;

use bytes::BytesMut;
use chrono::Utc;
use dashmap::DashMap;
use futures::{SinkExt, StreamExt};
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response};
use tokio_tungstenite::tungstenite::http::StatusCode;
use tokio_tungstenite::tungstenite::protocol::Message;

use futu_auth::{KeyRecord, KeyStore, RuntimeCounters, Scope};
use futu_auth_pipeline::{
    AuthDecision, AuthEnvelope, Credential, Endpoint, FilterRegistry, RejectKind, SurfaceId,
    authenticate_request,
};

/// v1.4.106 D1 5c: WS surface adapter — `AuthDecision::Reject` 翻成 silent drop.
///
/// **历史**: WS 在 v1.4.103 / v1.4.104 阶段 3 都按 silent drop 处理 reject
/// (与 v1.4.103 行为一致, 防 timing 探测 — 给客户端任何 wire response 都让
/// 它知道帧被读了 vs 没读). v1.4.106 D1 把这层"翻译为 unit"也走 trait, 让 4
/// surface SurfaceAdapter 一致.
///
/// **WireResponse = ()**: WS 不发任何东西回 client; 调用方拿 `Option<()>`
/// 知道是 reject (Some) 还是 allow (None) 即可继续 dispatch / 丢弃.
///
/// **不变量**: `translate_reject` 不能写日志 (pipeline 已 audit::reject 一次,
/// 不要重复). 只 drain reason / kind 让它进 `_`.
pub struct WsAdapter;

impl futu_auth_pipeline::SurfaceAdapter for WsAdapter {
    type WireResponse = ();

    fn surface_id() -> SurfaceId {
        SurfaceId::Ws
    }

    fn translate_reject(kind: RejectKind, reason: String) -> Self::WireResponse {
        // 与 v1.4.103/104 行为一致: silent drop. pipeline 已 audit reject,
        // 不再写 log; reason/kind drained 防 client 探测 daemon 状态.
        let _ = (kind, reason);
    }
}
use futu_codec::frame::FutuFrame;
use futu_codec::header::{FutuHeader, HEADER_SIZE, ProtoFmtType};
use futu_core::proto_id;

use crate::conn::{ClientConn, ConnState, DisconnectNotify, IncomingRequest};
use crate::listener::{MAX_CONNECTIONS, ServerConfig};
use crate::router::RequestRouter;

/// WebSocket 服务端
pub struct WsServer {
    listen_addr: String,
    config: ServerConfig,
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
    /// v1.0：握手时做 Bearer token 鉴权。None 或 `!is_configured()` → legacy 模式放行
    key_store: Option<Arc<KeyStore>>,
    /// v1.0：跨 REST / gRPC / WS 共享的限额 counters
    counters: Option<Arc<RuntimeCounters>>,
    /// v1.4.104 阶段 3: 跨 surface 共享的 response filter registry (proto 2001
    /// TRD_GET_ACC_LIST 默认装入). None 时 fallback 到内置 with_defaults.
    filter_registry: Option<Arc<FilterRegistry>>,
}

/// Shared dependencies for the WebSocket server.
///
/// These are the same runtime objects the raw TCP server owns. Keeping them in a
/// bundle avoids every constructor carrying a growing positional list as auth /
/// counters / filters evolve.
pub struct WsServerDeps {
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
}

impl WsServerDeps {
    pub fn new(
        connections: Arc<DashMap<u64, ClientConn>>,
        router: Arc<RequestRouter>,
        subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
    ) -> Self {
        Self {
            connections,
            router,
            subscriptions,
        }
    }
}

impl WsServer {
    /// 创建 WsServer，共享 TCP 的连接池、路由器、订阅管理器（无鉴权，向后兼容）
    pub fn new(
        listen_addr: String,
        config: ServerConfig,
        connections: Arc<DashMap<u64, ClientConn>>,
        router: Arc<RequestRouter>,
        subscriptions: Option<Arc<crate::subscription::SubscriptionManager>>,
    ) -> Self {
        Self::with_auth(
            listen_addr,
            config,
            WsServerDeps::new(connections, router, subscriptions),
            None,
            None,
        )
    }

    /// v1.0 入口：同时接入 KeyStore + 共享 RuntimeCounters 做握手鉴权和 per-message
    /// scope / 限额检查。`key_store = None` 或未配置时保持 legacy（全放行）。
    pub fn with_auth(
        listen_addr: String,
        config: ServerConfig,
        deps: WsServerDeps,
        key_store: Option<Arc<KeyStore>>,
        counters: Option<Arc<RuntimeCounters>>,
    ) -> Self {
        Self {
            listen_addr,
            config,
            connections: deps.connections,
            router: deps.router,
            subscriptions: deps.subscriptions,
            key_store,
            counters,
            filter_registry: None,
        }
    }

    /// v1.4.104 阶段 3: 注入显式 FilterRegistry (跨 surface 共享同一份).
    /// 不调用此 setter 则 run() 时 fallback 到 `FilterRegistry::with_defaults()`.
    pub fn with_filter_registry(mut self, registry: Arc<FilterRegistry>) -> Self {
        self.filter_registry = Some(registry);
        self
    }

    /// 启动 WebSocket 服务端监听
    pub async fn run(&self) -> anyhow::Result<()> {
        let listener = TcpListener::bind(&self.listen_addr).await?;
        tracing::info!(addr = %self.listen_addr, "WebSocket server listening");

        let (req_tx, req_rx) = mpsc::unbounded_channel::<IncomingRequest>();
        let (disconnect_tx, mut disconnect_rx) = mpsc::unbounded_channel::<DisconnectNotify>();

        // 启动请求处理任务（与 TCP 共享同一逻辑）
        let connections = Arc::clone(&self.connections);
        let router = Arc::clone(&self.router);
        let config = self.config.clone();
        // v1.4.104 阶段 3: KeyStore + RuntimeCounters 总是材料化 (legacy mode 用
        // empty / new()), 让 ws_process_requests 拿 non-Option Arc 直接调 pipeline.
        // 行为与 v1.4.103 等价: KeyStore::empty().is_configured() = false, pipeline
        // 走 legacy short-circuit (Allow{rec:None} 不 audit, body-aware 不 enforce).
        let key_store_for_process = self
            .key_store
            .clone()
            .unwrap_or_else(|| Arc::new(KeyStore::empty()));
        let counters_for_process = self
            .counters
            .clone()
            .unwrap_or_else(|| Arc::new(RuntimeCounters::new()));
        let filter_registry_for_process = self
            .filter_registry
            .clone()
            .unwrap_or_else(|| Arc::new(FilterRegistry::with_defaults()));
        tokio::spawn(async move {
            ws_process_requests(
                req_rx,
                connections,
                router,
                config,
                counters_for_process,
                key_store_for_process,
                filter_registry_for_process,
            )
            .await;
        });

        // 启动连接清理任务
        let cleanup_connections = Arc::clone(&self.connections);
        let cleanup_subs = self.subscriptions.clone();
        tokio::spawn(async move {
            while let Some(notify) = disconnect_rx.recv().await {
                let removed = cleanup_connections.remove(&notify.conn_id);
                if removed.is_some() {
                    if let Some(ref subs) = cleanup_subs {
                        subs.on_disconnect(notify.conn_id);
                    }
                    tracing::info!(
                        conn_id = notify.conn_id,
                        remaining = cleanup_connections.len(),
                        "ws connection removed from pool"
                    );
                }
            }
        });

        // 接受连接循环
        let connections = Arc::clone(&self.connections);
        let key_store_accept = self.key_store.clone();
        // v1.4.104 阶段 3: scope_mode 局部计算 (KeyStore configured = 启用 auth).
        let scope_mode = self.key_store.as_ref().is_some_and(|ks| ks.is_configured());
        if !scope_mode {
            // v1.4.93 P0-5 (NEW-C-02): 加强 legacy mode loud WARN —
            // 对齐 REST mutating-blocked policy 的 startup signal。任何
            // 未授权客户端都可 handshake + 接收 push（HTTP 101 OK）。本版
            // **不 reject**（保持向后兼容），未来 v2 默认 reject。
            tracing::warn!("{}", legacy_mode_warn_tracing_message());
            eprintln!("{}", legacy_mode_warn_stderr_message());
        }
        loop {
            let (stream, peer_addr) = listener.accept().await?;

            if connections.len() >= MAX_CONNECTIONS {
                tracing::warn!(
                    peer = %peer_addr,
                    "max connections reached ({}), rejecting ws client",
                    MAX_CONNECTIONS,
                );
                drop(stream);
                continue;
            }

            let conn_id = ClientConn::generate_conn_id();
            let aes_key = ClientConn::generate_aes_key();
            stream.set_nodelay(true).ok();

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

            let (tx, authed) = run_ws_connection(
                stream,
                conn_id,
                aes_key,
                req_tx.clone(),
                disconnect_tx.clone(),
                key_store_accept.clone(),
            )
            .await;

            // 握手鉴权失败 → run_ws_connection 已经 drop 连接；这里什么都不做
            let Some(authed) = authed else {
                continue;
            };

            let (key_id, scopes, allowed_markets, allowed_acc_ids) = match authed {
                AuthResult::Authenticated(rec) => (
                    Some(rec.id.clone()),
                    rec.scopes.clone(),
                    // v1.4.105 D3 (Phase 4) T-B2: 拷贝 caller key 的 allowed_markets
                    // 到 ClientConn 让 PushDispatcher::push_trd_acc Layer 3 用.
                    rec.allowed_markets
                        .as_ref()
                        .map(|s| std::sync::Arc::new(s.clone())),
                    // codex round 1 F4 (P2) v1.4.105: 拷贝 caller key 的
                    // allowed_acc_ids 到 ClientConn 让 PushDispatcher::push_trd_acc
                    // Layer 1 push-time 硬过滤. 防 stale subscription /
                    // KeyRecord reload 后 acc 范围窄化 时 push leak.
                    rec.allowed_acc_ids
                        .as_ref()
                        .map(|s| std::sync::Arc::new(s.clone())),
                ),
                AuthResult::Legacy => (None, HashSet::new(), None, None),
            };

            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,
                keepalive_count: std::sync::atomic::AtomicU32::new(0),
                tx,
                key_id,
                scopes,
                allowed_markets,
                allowed_acc_ids,
            };

            connections.insert(conn_id, conn);
        }
    }
}

/// 握手鉴权结果：scope 模式下的 KeyRecord 或 legacy 全放行
enum AuthResult {
    Authenticated(Arc<KeyRecord>),
    Legacy,
}

// `ErrorResponse` 是 tungstenite/axum 的 ws upgrade 接口约定类型，含
// `http::Response<Option<String>>` 必然偏大（≥136 bytes）。本 fn 在 ws
// 握手 error path（slot poisoning，极罕见），非 hot path；改 `Box<ErrorResponse>`
// 需要所有 caller deref 适配，收益不抵成本。按 pitfall #51"对齐 C++ = 减法"
// 保最小侵入。
#[allow(clippy::result_large_err)]
fn store_ws_auth_result(
    slot: &std::sync::Mutex<Option<AuthResult>>,
    result: AuthResult,
    conn_id: u64,
) -> Result<(), ErrorResponse> {
    match slot.lock() {
        Ok(mut guard) => {
            *guard = Some(result);
            Ok(())
        }
        Err(e) => {
            tracing::warn!(conn_id = conn_id, error = %e, "ws auth slot poisoned during handshake");
            Err(make_err_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                "ws auth state unavailable",
            ))
        }
    }
}

fn take_ws_auth_result(
    slot: &std::sync::Mutex<Option<AuthResult>>,
    conn_id: u64,
) -> Option<AuthResult> {
    match slot.lock() {
        Ok(mut guard) => match guard.take() {
            Some(result) => Some(result),
            None => {
                tracing::warn!(
                    conn_id = conn_id,
                    "ws handshake succeeded without auth state; closing connection"
                );
                None
            }
        },
        Err(e) => {
            tracing::warn!(conn_id = conn_id, error = %e, "ws auth slot poisoned after handshake");
            None
        }
    }
}

/// 运行单个 WebSocket 连接的收发循环
///
/// 接收端：WS Binary Message → 解析为 FutuFrame → 发到 req_tx
/// 发送端：frame_rx 收到 FutuFrame → 编码为字节 → 发送 WS Binary Message
///
/// 返回 `(frame_tx, Option<AuthResult>)` —— `None` 表示握手 / 鉴权失败，调用方
/// 不应把这个 conn_id 加入连接池。
async fn run_ws_connection(
    stream: tokio::net::TcpStream,
    conn_id: u64,
    _aes_key: [u8; 16],
    req_tx: mpsc::UnboundedSender<IncomingRequest>,
    disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
    key_store: Option<Arc<KeyStore>>,
) -> (mpsc::Sender<FutuFrame>, Option<AuthResult>) {
    let (frame_tx, mut frame_rx) = mpsc::channel::<FutuFrame>(256);

    // 握手回调里验证 token + scope，结果存到 authed_slot；accept_hdr_async 内部
    // 只读 callback 的 Ok/Err 决定要不要升级，所以 KeyRecord 要借 Mutex 传出来
    let authed_slot: Arc<std::sync::Mutex<Option<AuthResult>>> =
        Arc::new(std::sync::Mutex::new(None));
    let slot_cb = Arc::clone(&authed_slot);
    let store_cb = key_store.clone();

    #[allow(clippy::result_large_err)] // ErrorResponse 是 tungstenite 的类型，我们无法改其大小
    let callback = move |req: &Request, resp: Response| -> Result<Response, ErrorResponse> {
        // v1.4.102 F-002 fix (P2, leaf v1.4.100 报告): WS handshake 加 Origin 校验.
        // 历史: legacy WS 看到任意 evil Origin 仍 handshake 成功 (HTTP 101).
        // 浏览器 WS exfiltration POC 还没证实, 但默认行为该收紧.
        //
        // **策略** (与 REST CORS BUG-008 对齐):
        //   1. `FUTU_WS_ALLOWED_ORIGINS=https://app.com,http://localhost:3000`
        //      显式 allowlist (推荐).
        //   2. key_store 配置了 (auth enabled) + 未设 env → loopback only.
        //   3. legacy unauth + 未设 env → loopback only (与 REST CORS 对齐).
        //   4. **没 Origin 头**(非 browser 客户端 / native WS lib)→ 通过.
        if let Some(origin_hv) = req.headers().get("origin") {
            let origin_str = match origin_hv.to_str() {
                Ok(s) => s,
                Err(e) => {
                    tracing::warn!(conn_id, error = %e, "ws Origin header is not valid UTF-8");
                    futu_auth::audit::reject("ws", "/ws", "<origin>", "invalid Origin header");
                    return Err(make_err_response(
                        StatusCode::FORBIDDEN,
                        "Invalid Origin header",
                    ));
                }
            };
            let allowed = ws_check_origin(origin_str, store_cb.as_deref());
            if !allowed {
                futu_auth::audit::reject("ws", "/ws", "<origin>", "Origin rejected by allowlist");
                return Err(make_err_response(
                    StatusCode::FORBIDDEN,
                    "Origin not allowed (configure FUTU_WS_ALLOWED_ORIGINS env)",
                ));
            }
        }

        // legacy：未注入 KeyStore 或未配置 keys.json → 全放行
        let Some(store) = store_cb.as_ref() else {
            store_ws_auth_result(slot_cb.as_ref(), AuthResult::Legacy, conn_id)?;
            return Ok(resp);
        };
        if !store.is_configured() {
            store_ws_auth_result(slot_cb.as_ref(), AuthResult::Legacy, conn_id)?;
            return Ok(resp);
        }

        // 提取 token：?token=... 或 Authorization: Bearer ...
        let token = extract_ws_token(req);
        let Some(token) = token else {
            futu_auth::audit::reject("ws", "/ws", "<missing>", "missing token");
            return Err(make_err_response(
                StatusCode::UNAUTHORIZED,
                "missing api key (use ?token=... or Authorization: Bearer ...)",
            ));
        };

        let Some(rec) = store.verify(&token) else {
            futu_auth::audit::reject("ws", "/ws", "<invalid>", "invalid api key");
            return Err(make_err_response(
                StatusCode::UNAUTHORIZED,
                "invalid api key",
            ));
        };

        if rec.is_expired(Utc::now()) {
            futu_auth::audit::reject("ws", "/ws", &rec.id, "key expired");
            return Err(make_err_response(StatusCode::UNAUTHORIZED, "key expired"));
        }

        // 最低门槛：qot:read。真正按 proto_id 的细粒度检查留给 process_requests
        // v1.4.102 codex 24 F4 (P2): forbidden body 通用化, 不再泄 scope.
        // 详细 audit 信息留 audit log; HTTP body 用通用 "forbidden". REST + gRPC
        // + WS handshake 同步 (BUG-011 + 24 F4).
        if !rec.scopes.contains(&Scope::QotRead) {
            futu_auth::audit::reject("ws", "/ws", &rec.id, "missing qot:read");
            return Err(make_err_response(StatusCode::FORBIDDEN, "forbidden"));
        }

        futu_auth::audit::allow("ws", "/ws", &rec.id, Some("qot:read"));
        store_ws_auth_result(slot_cb.as_ref(), AuthResult::Authenticated(rec), conn_id)?;
        Ok(resp)
    };

    let ws_stream = match tokio_tungstenite::accept_hdr_async(stream, callback).await {
        Ok(ws) => ws,
        Err(e) => {
            tracing::warn!(conn_id = conn_id, error = %e, "ws handshake failed");
            let _ = disconnect_tx.send(DisconnectNotify { conn_id });
            return (frame_tx, None);
        }
    };

    // 握手成功 → slot_cb 应已填；若内部状态不可用，fail-closed 关闭这条 WS。
    let Some(authed) = take_ws_auth_result(authed_slot.as_ref(), conn_id) else {
        let _ = disconnect_tx.send(DisconnectNotify { conn_id });
        return (frame_tx, None);
    };

    let (mut ws_sink, mut ws_stream_rx) = ws_stream.split();

    // 发送任务：FutuFrame → 编码为字节 → WS Binary
    tokio::spawn(async move {
        while let Some(frame) = frame_rx.recv().await {
            let mut buf = BytesMut::new();
            frame.header.encode(&mut buf);
            buf.extend_from_slice(&frame.body);
            // tokio-tungstenite 0.26+: Message::Binary 收 Bytes（之前是 Vec<u8>）
            let msg = Message::Binary(buf.freeze());
            if let Err(e) = ws_sink.send(msg).await {
                tracing::warn!(conn_id = conn_id, error = %e, "ws send failed");
                break;
            }
        }
    });

    // 接收任务：WS Binary → 解析 FutuFrame → req_tx
    tokio::spawn(async move {
        while let Some(result) = ws_stream_rx.next().await {
            match result {
                Ok(msg) => {
                    let data = match msg {
                        Message::Binary(data) => data,
                        Message::Close(_) => {
                            tracing::info!(conn_id = conn_id, "ws client sent close");
                            break;
                        }
                        Message::Ping(_) | Message::Pong(_) => {
                            // tokio-tungstenite 自动处理 ping/pong
                            continue;
                        }
                        _ => {
                            // 忽略 Text 等其他类型
                            continue;
                        }
                    };

                    // 解析 FutuAPI 帧（44 字节帧头 + body）
                    if data.len() < HEADER_SIZE {
                        tracing::warn!(
                            conn_id = conn_id,
                            len = data.len(),
                            "ws message too short for futu header"
                        );
                        continue;
                    }

                    let header_buf = BytesMut::from(&data[..]);
                    let header = match FutuHeader::peek(&header_buf) {
                        Ok(Some(h)) => h,
                        Ok(None) => {
                            tracing::warn!(conn_id = conn_id, "ws header peek returned None");
                            continue;
                        }
                        Err(e) => {
                            tracing::warn!(conn_id = conn_id, error = %e, "ws invalid futu header");
                            continue;
                        }
                    };

                    let expected_len = HEADER_SIZE + header.body_len as usize;
                    if data.len() < expected_len {
                        tracing::warn!(
                            conn_id = conn_id,
                            expected = expected_len,
                            actual = data.len(),
                            "ws message shorter than expected frame size"
                        );
                        continue;
                    }

                    let body = bytes::Bytes::copy_from_slice(&data[HEADER_SIZE..expected_len]);

                    let req = IncomingRequest::builder(
                        conn_id,
                        header.proto_id,
                        header.serial_no,
                        header.proto_fmt_type,
                        body,
                    )
                    .build();

                    if req_tx.send(req).is_err() {
                        break;
                    }
                }
                Err(e) => {
                    tracing::warn!(conn_id = conn_id, error = %e, "ws recv error");
                    break;
                }
            }
        }
        tracing::info!(conn_id = conn_id, "ws connection closed");
        let _ = disconnect_tx.send(DisconnectNotify { conn_id });
    });

    (frame_tx, Some(authed))
}

/// v1.4.93 P0-5 (NEW-C-02): legacy mode 的 `tracing::warn!` 内容。
///
/// 抽出 const fn 以便单测验证 warn 消息携带 "v2"/"reject" 等关键提示词，
/// 防止后续被误删（同模式 v1.4.86 SEC-003 Q4 已沉淀）。
pub(crate) const fn legacy_mode_warn_tracing_message() -> &'static str {
    "WS server running WITHOUT API key auth (legacy mode); \
     all WS clients accept unauthenticated handshake (no-token / \
     wrong-bearer / bogus-query all return success). \
     Pass KeyStore via with_auth() to enable. \
     v2 will default-reject; migrate to --rest-keys-file / --ws-keys-file for production."
}

/// v1.4.93 P0-5 (NEW-C-02): legacy mode 的 stderr 用户可见消息。
///
/// 比 tracing::warn 更短，方便 systemd / docker logs 一行抓住。
pub(crate) const fn legacy_mode_warn_stderr_message() -> &'static str {
    "⚠️  WS server (legacy mode, no --ws-keys-file): \
     unauthenticated handshakes accepted. v2 will default-reject. \
     Migrate to --ws-keys-file for production."
}

/// 从握手 Request 里提取 token:优先 ?token=...,再 Authorization: Bearer ...
///
/// 浏览器 WS API 不允许设置自定义 header,所以优先支持 query;原生客户端
/// (curl / Futu SDK 之类)两种都可以.
///
/// v1.4.104 阶段 7-3 fix:Authorization 走 `parse_bearer_scheme` 共享 helper
/// (case-insensitive RFC 7235). 旧实装 `strip_prefix("Bearer ")` 是
/// case-sensitive bug — `bearer xxx` / `BEARER xxx` 会被拒, 与 gRPC + REST +
/// MCP 行为不一致 (它们已 case-insensitive). 本 commit 修齐.
fn extract_ws_token(req: &Request) -> Option<String> {
    if let Some(q) = req.uri().query() {
        // 手写 query 解析,避免引 url / percent-encoding 新依赖
        let params: HashMap<&str, &str> =
            q.split('&').filter_map(|kv| kv.split_once('=')).collect();
        if let Some(v) = params.get("token")
            && !v.is_empty()
        {
            return Some((*v).to_string());
        }
    }
    let header = req.headers().get("authorization")?;
    let value = header.to_str().ok()?;
    futu_auth_pipeline::parse_bearer_scheme(value).map(|t| t.to_string())
}

/// v1.4.102 F-002 fix (P2): WS handshake Origin 校验.
///
/// **匹配规则** (与 REST CORS BUG-008 对齐):
///   1. `FUTU_WS_ALLOWED_ORIGINS` env 显式 allowlist → 严格匹配
///   2. 未设 env + key_store configured → loopback (127.0.0.1 / localhost / [::1])
///   3. 未设 env + legacy unauth → 允许任意 (向后兼容)
///
/// 浏览器以外的原生 WS 客户端 (futucli, Python SDK 等) **不发 Origin 头**,
/// 调用方应在 callback 里**先 check Origin header presence**, 缺失时跳过本 fn.
fn ws_check_origin(origin: &str, key_store: Option<&KeyStore>) -> bool {
    if let Ok(raw) = std::env::var("FUTU_WS_ALLOWED_ORIGINS") {
        let trimmed = raw.trim();
        if !trimmed.is_empty() {
            for allowed in trimmed
                .split(',')
                .map(|s| s.trim())
                .filter(|s| !s.is_empty())
            {
                if allowed == origin {
                    return true;
                }
            }
            // env 显式设了但 origin 不在 list → reject
            return false;
        }
    }
    // env 未设 / 空: 按 auth 状态决定 default
    let auth_enabled = key_store.is_some_and(|ks| ks.is_configured());
    if !auth_enabled {
        // v1.4.104 eli P2-001 (P2) fix: legacy unauth WS 也 default loopback,
        // 不再允许任意 Origin. 之前 evil 站点可发 WS handshake 拿 101, 配合
        // socket 抓数据 → cross-origin leak 风险. 用户要 cross-origin browser
        // client 必须显式 FUTU_WS_ALLOWED_ORIGINS=https://app.example.com.
        // (与 REST CORS v1.4.104 P1-004 对齐.)
        return is_strict_loopback_origin(origin);
    }
    // auth enabled: loopback only.
    // v1.4.102 codex 25 F5 / 30 F3 / 31 F6 (P2/P3) refine: 严格 host 匹配,
    // 不接受 prefix-混入 evil hosts (`[::1].evil.com` / `localhost.evil.com`).
    is_strict_loopback_origin(origin)
}

/// v1.4.102 codex 25 F5 / 30 F3 / 31 F6 strict parser. 不接受 path / query /
/// userinfo / 任何混入字符. 与 `crates/futu-rest/src/server.rs::is_loopback_origin_str`
/// 行为对齐 (REST + WS handshake 共用 same strict policy).
fn is_strict_loopback_origin(s: &str) -> bool {
    let after_scheme = match s
        .strip_prefix("http://")
        .or_else(|| s.strip_prefix("https://"))
    {
        Some(rest) => rest,
        None => return false,
    };
    if after_scheme.contains('/')
        || after_scheme.contains('?')
        || after_scheme.contains('#')
        || after_scheme.contains('@')
    {
        return false;
    }
    // v1.4.102 codex 41 F4: 区分 None vs Some("") (host: vs host).
    let (host, port_opt): (&str, Option<&str>) =
        if let Some(rest) = after_scheme.strip_prefix("[::1]") {
            if rest.is_empty() {
                ("[::1]", None)
            } else if let Some(p) = rest.strip_prefix(':') {
                ("[::1]", Some(p))
            } else {
                return false;
            }
        } else if let Some((h, p)) = after_scheme.rsplit_once(':') {
            (h, Some(p))
        } else {
            (after_scheme, None)
        };
    if !matches!(host, "127.0.0.1" | "localhost" | "[::1]") {
        return false;
    }
    if let Some(port_str) = port_opt {
        match port_str.parse::<u16>() {
            Ok(p) if p >= 1 => {}
            _ => return false,
        }
    }
    true
}

/// 构造 tungstenite 握手失败的 HTTP 响应
fn make_err_response(code: StatusCode, msg: &str) -> ErrorResponse {
    let body = Some(format!(r#"{{"error":"{msg}"}}"#));
    let mut resp = tokio_tungstenite::tungstenite::http::Response::new(body);
    *resp.status_mut() = code;
    resp.headers_mut().insert(
        "content-type",
        tokio_tungstenite::tungstenite::http::HeaderValue::from_static("application/json"),
    );
    resp
}

/// 处理 WebSocket 连接的请求（逻辑与 TCP 的 process_requests 相同，额外做 scope / 限额）
///
/// v1.4.104 阶段 3 重构: inline scope gate + trade:real rate gate +
/// `ws_body_aware_check` 三段折叠为单一 `authenticate_request` pipeline 调用.
///
/// **流程**:
/// 1. AES decrypt (handshake INIT_CONNECT 跳过 — body 是明文 RSA 加密 InitReq)
/// 2. 非 INIT_CONNECT + scope_for_proto_id != None → 调 pipeline:
///    - Credential::PreVerified(`KeyStore::get_by_id(key_id)`) 拿 SIGHUP-aware
///      最新 rec, 复用 handshake 已 verify 的身份 (skip re-verify).
///    - Endpoint::Proto(proto_id), commit_rate=true (per-msg rate 闸门).
///    - Reject → drop request (audit 已 emit). Allow → 拿 allowed_acc_ids.
/// 3. dispatch (router / handle_init_connect / handle_keepalive)
/// 4. response filter (proto 2001 等, 通过 FilterRegistry::apply).
///
/// **legacy mode 行为**: KeyStore::empty().is_configured() = false, pipeline
/// 走 legacy short-circuit (Allow{rec:None}, 不 audit, body-aware 不 enforce).
/// 1xxx 系统协议 (INIT_CONNECT / KEEP_ALIVE / GET_GLOBAL_STATE) 与 v1.4.103
/// 一致跳过 pipeline 直接 dispatch (handshake / heartbeat / 公开协议无 scope).
async fn ws_process_requests(
    mut req_rx: mpsc::UnboundedReceiver<IncomingRequest>,
    connections: Arc<DashMap<u64, ClientConn>>,
    router: Arc<RequestRouter>,
    config: ServerConfig,
    counters: Arc<RuntimeCounters>,
    key_store: Arc<KeyStore>,
    filter_registry: Arc<FilterRegistry>,
) {
    use crate::listener::ApiServer;

    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;

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

        // ── Step 0: v1.4.106 codex 0532 F3 (P2): daemon-internal proto_id
        // (高位 0x8000_0000 bit) 绝不应从 raw WS 公开 surface 进入 — 仅 REST
        // handler 内部合成给 router. 在 AES decrypt 前显式 reject + log,
        // 防探测 daemon 内部 routing.
        if futu_auth::is_internal_proto_id(proto_id_val) {
            tracing::warn!(
                conn_id,
                proto_id = proto_id_val,
                "rejecting daemon-internal proto_id at raw WS public surface (codex 0532 F3)"
            );
            continue;
        }

        // ── Step 1: AES decrypt (always for non-INIT_CONNECT) ─────────────────
        // pipeline + body-aware 都需 plaintext body. INIT_CONNECT body 是 RSA
        // 加密的 ConnInitReq, 由 handle_init_connect 自行 RSA decrypt.
        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) => {
                    tracing::warn!(
                        conn_id = conn_id,
                        proto_id = proto_id_val,
                        error = %e,
                        "ws AES decrypt request failed, dropping"
                    );
                    continue;
                }
            }
        }

        // ── Step 2: Pipeline auth ─────────────────────────────────────────────
        // INIT_CONNECT (handshake) + 1xxx 系统协议 (scope_for_proto_id == None)
        // 跳 pipeline. 与 v1.4.103 行为一致: 系统协议不审 / 不限频 / 直 dispatch.
        let needed_scope = futu_auth_pipeline::capability::scope_for_proto_id(proto_id_val);
        // codex 0522 F1 v1.4.106: 提前抓 conn.key_id 快照让 dispatch IncomingRequest
        // 也能填 caller_key_id (per-call snapshot, 与 caller_allowed_acc_ids 同源).
        // 即使 INIT_CONNECT / 1xxx 系统协议 (跳 pipeline) 也带上 — handler 可基于
        // key_id 做 per-key 订阅配额 / cleanup / 审计.
        let dispatch_caller_key_id: Option<String> =
            connections.get(&conn_id).and_then(|c| c.key_id.clone());
        let allowed_acc_ids_for_resp_filter: Option<HashSet<u64>> =
            if proto_id_val == proto_id::INIT_CONNECT || needed_scope.is_none() {
                None
            } else {
                // 从 conn 取 key_id 快照, 然后 KeyStore::get_by_id 拿 SIGHUP-aware
                // 最新 rec (limits / scopes / allowed_acc_ids 都跟最新). 找不到 ↦
                // Credential::None (legacy mode 放行 / scope mode reject Unauth).
                let key_id_snap = dispatch_caller_key_id.clone();
                let rec_opt = key_id_snap.as_ref().and_then(|id| key_store.get_by_id(id));
                let credential = match rec_opt {
                    Some(rec) => Credential::PreVerified(rec),
                    None => Credential::None,
                };

                let env = AuthEnvelope {
                    surface: SurfaceId::Ws,
                    endpoint: Endpoint::Proto(proto_id_val),
                    needed_scope,
                    credential,
                    proto_id: Some(proto_id_val),
                    body: &req.body,
                    explicit_acc_id: None,
                    explicit_ctx: None,
                    commit_rate: true, // WS per-msg 是 trade:real 唯一 rate gate
                    audit_emit: true,
                };

                // v1.4.106 D1 5c: 走 SurfaceAdapter trait
                // (`WsAdapter::translate_decision`), 与 4 surface 一致.
                // Allow → Some(allowed_acc_ids), Reject → None + silent drop.
                use futu_auth_pipeline::SurfaceAdapter;
                match authenticate_request(&key_store, &counters, env) {
                    AuthDecision::Allow {
                        allowed_acc_ids, ..
                    } => allowed_acc_ids,
                    decision @ AuthDecision::Reject { .. } => {
                        // WsAdapter::translate_decision 返 Some(()) 表 reject
                        // (silent drop). pipeline 已 audit reject, 这里不打 log.
                        let _ = WsAdapter::translate_decision(decision);
                        continue;
                    }
                }
            };

        // ── Step 3: Dispatch ──────────────────────────────────────────────────
        let response_body = match proto_id_val {
            proto_id::INIT_CONNECT => match connections.get_mut(&conn_id) {
                Some(mut conn) => conn
                    .handle_init_connect(
                        &req.body,
                        config.server_ver,
                        config.login_user_id,
                        config.keepalive_interval,
                        config.rsa_private_key.as_deref(),
                    )
                    .ok(),
                _ => None,
            },
            proto_id::KEEP_ALIVE => match connections.get(&conn_id) {
                Some(conn) => conn.handle_keepalive(&req.body).ok(),
                _ => None,
            },
            _ => {
                // v1.4.105 D2 T-A1 fix: caller_allowed_acc_ids 从 pipeline allow
                // decision 真填进 IncomingRequest, 让 dispatch handler (e.g.
                // SubAccPushHandler) 端 enforce per-acc whitelist defense-in-depth.
                // codex 0522 F1 v1.4.106: 同步填 caller_key_id (per-call snapshot
                // 来自 conn.key_id), 让 cross-surface handler 都能识别 caller.
                let dispatch_req = IncomingRequest::builder(
                    req.conn_id,
                    req.proto_id,
                    req.serial_no,
                    req.proto_fmt_type,
                    req.body.clone(),
                )
                .with_idempotency_key(req.idempotency_key.clone())
                .with_caller_scope(
                    allowed_acc_ids_for_resp_filter
                        .as_ref()
                        .map(|s| std::sync::Arc::new(s.clone())),
                    dispatch_caller_key_id.clone(),
                )
                .build();
                router.dispatch(conn_id, &dispatch_req).await
            }
        };

        // ── Step 4: Response filter (TRD_GET_ACC_LIST 等) ────────────────────
        if let Some(body) = response_body {
            // FilterRegistry::apply(): proto_id 未注册 → 原 body 不动 (no-op).
            // 注册了 (e.g. proto 2001) → filter by allowed_acc_ids.
            let filtered =
                filter_registry.apply(proto_id_val, body, allowed_acc_ids_for_resp_filter.as_ref());
            ApiServer::send_response(&connections, conn_id, proto_id_val, serial_no, filtered)
                .await;
        }
    }
}

#[cfg(test)]
mod tests;