futu_rest/

ws.rs

//! WebSocket 推送模块
//!
//! 在 REST API 端口上提供 /ws 路由，客户端通过 WebSocket 接收实时推送。
//!
//! 推送事件通过 broadcast channel 从 OpenD 核心分发到所有 WebSocket 客户端。

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

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
use chrono::Utc;
use futures::{SinkExt, StreamExt};
use tokio::sync::broadcast;

use futu_auth::{KeyRecord, KeyStore, Scope};
use futu_server::push::ExternalPushSink;

use crate::adapter::RestState;

/// WebSocket 推送事件
#[derive(Clone, Debug, serde::Serialize)]
pub struct WsPushEvent {
    /// 推送类型: "quote", "trade", "notify"
    #[serde(rename = "type")]
    pub event_type: String,
    /// 该事件需要哪个 scope 才能被某个 client 接收（filter 用，不发到客户端）
    #[serde(skip)]
    pub required_scope: WsPushScope,
    /// 协议 ID
    pub proto_id: u32,
    /// 证券标识 (行情推送)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sec_key: Option<String>,
    /// 订阅类型 (行情推送)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sub_type: Option<i32>,
    /// **v1.4.106 codex 1131 F4 [P1]**: rehab 类型 (KL push 非 0, 其它 sub_type
    /// 为 0). 客户端用于 (sec_key, sub_type, rehab_type) 三元 key 自行 filter
    /// 不感兴趣的 KL rehab 推送.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rehab_type: Option<i32>,
    /// 交易账户 ID (交易推送)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub acc_id: Option<u64>,
    /// protobuf body 的 base64 编码
    pub body_b64: String,
    /// v1.4.105 D3 (Phase 4) T-B1: 交易推送的 trd_market 大写字符串 ("HK" /
    /// "US" / "CN" / "HKCC" / "FUTURES" / "SG" / "AU" / "JP" / "MY" / "CA").
    /// PushDispatcher 一次 decode body 后透传过来, 让 WS push filter Layer 3
    /// (allowed_markets) 直接读. `None` = 非 trade event / decode 失败 /
    /// market 未知 (Layer 3 向后兼容不 trigger drop).
    ///
    /// 客户端可见: trade event 出现 `trd_market` 字段, qot/notify 不出现
    /// (`skip_serializing_if = "Option::is_none"`).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub trd_market: Option<String>,
}

/// WS 推送事件需要的最低 scope（client 没这个 scope 就收不到）
///
/// - `Quote` → `qot:read`：行情类
/// - `Notify` → `qot:read`：通用通知（如订阅状态、网关心跳）
/// - `Trade` → `acc:read`：交易回报涉及账户隐私，必须有账户读权限
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum WsPushScope {
    /// 行情推送（订阅 symbol 后 push 的 basic_qot / order_book / ticker 等）。
    /// 需要 [`Scope::QotRead`]。
    #[default]
    Quote,
    /// 广播通知（系统事件 / 全局消息）。需要 [`Scope::QotRead`]。
    Notify,
    /// 交易推送（订单状态变化 / 成交回报）。需要 [`Scope::AccRead`]。
    Trade,
}

impl WsPushScope {
    /// 该事件类型需要的 Scope；client 必须持有这个 scope 才能收到
    pub fn required_scope(&self) -> Scope {
        match self {
            WsPushScope::Quote => Scope::QotRead,
            WsPushScope::Notify => Scope::QotRead,
            WsPushScope::Trade => Scope::AccRead,
        }
    }
}

/// WebSocket 推送广播器
///
/// OpenD 核心推送事件 → broadcast channel → 所有 WebSocket 客户端
///
/// 实现 `ExternalPushSink` trait，可直接嵌入 PushDispatcher。
#[derive(Clone)]
pub struct WsBroadcaster {
    tx: broadcast::Sender<WsPushEvent>,
}

impl WsBroadcaster {
    pub fn new(capacity: usize) -> Self {
        let (tx, _) = broadcast::channel(capacity);
        Self { tx }
    }

    /// 发送推送事件到所有 WebSocket 客户端
    pub fn send(&self, event: WsPushEvent) {
        // 忽略没有接收者的情况
        let _ = self.tx.send(event);
    }

    /// 创建接收端
    pub fn subscribe(&self) -> broadcast::Receiver<WsPushEvent> {
        self.tx.subscribe()
    }

    fn encode_body(body: &[u8]) -> String {
        use base64::Engine;
        base64::engine::general_purpose::STANDARD.encode(body)
    }

    /// 发送行情推送.
    ///
    /// **v1.4.106 codex 1131 F4 [P1]**: 加 `rehab_type` 参数. KL push 的
    /// `rehab_type` ≠ 0, 其它 sub_type → 0. 当前 REST WS 仍 broadcast 所有
    /// quote events 给 qot:read 订阅者 (per-conn 三元 key 过滤是 raw TCP 专属
    /// 行为 — REST WS 用 broadcast 模型). 但 rehab_type 透传给客户端可见, 让
    /// agent 自己识别 KL push 的 rehab 类型.
    pub fn push_quote(
        &self,
        sec_key: &str,
        sub_type: i32,
        rehab_type: i32,
        proto_id: u32,
        body: &[u8],
    ) {
        self.send(WsPushEvent {
            event_type: "quote".to_string(),
            required_scope: WsPushScope::Quote,
            proto_id,
            sec_key: Some(sec_key.to_string()),
            sub_type: Some(sub_type),
            rehab_type: Some(rehab_type),
            acc_id: None,
            body_b64: Self::encode_body(body),
            trd_market: None,
        });
    }

    /// 发送广播推送
    pub fn push_broadcast(&self, proto_id: u32, body: &[u8]) {
        self.send(WsPushEvent {
            event_type: "notify".to_string(),
            required_scope: WsPushScope::Notify,
            proto_id,
            sec_key: None,
            sub_type: None,
            rehab_type: None,
            acc_id: None,
            body_b64: Self::encode_body(body),
            trd_market: None,
        });
    }

    /// 发送交易推送
    ///
    /// v1.4.105 D3 (Phase 4) T-B1: `trd_market` 由 [`PushDispatcher`] 一次
    /// decode body 后透传, 直接塞 [`WsPushEvent.trd_market`] 给后续 Layer 3
    /// filter 与客户端可见.
    pub fn push_trade(&self, acc_id: u64, proto_id: u32, body: &[u8], trd_market: Option<&str>) {
        self.send(WsPushEvent {
            event_type: "trade".to_string(),
            required_scope: WsPushScope::Trade,
            proto_id,
            sec_key: None,
            sub_type: None,
            rehab_type: None,
            acc_id: Some(acc_id),
            body_b64: Self::encode_body(body),
            trd_market: trd_market.map(|s| s.to_string()),
        });
    }
}

/// 实现 ExternalPushSink，使 WsBroadcaster 可嵌入 PushDispatcher
impl ExternalPushSink for WsBroadcaster {
    fn on_quote_push(
        &self,
        sec_key: &str,
        sub_type: i32,
        rehab_type: i32,
        proto_id: u32,
        body: &[u8],
    ) {
        self.push_quote(sec_key, sub_type, rehab_type, proto_id, body);
    }

    fn on_broadcast_push(&self, proto_id: u32, body: &[u8]) {
        self.push_broadcast(proto_id, body);
    }

    fn on_trade_push(&self, acc_id: u64, proto_id: u32, body: &[u8], trd_market: Option<&str>) {
        self.push_trade(acc_id, proto_id, body, trd_market);
    }
}

/// WebSocket 握手鉴权：从 `?token=xxx` 查询参数或 `Authorization: Bearer` header 提取 token
///
/// 浏览器 WebSocket API 不允许设置自定义 header，所以优先支持 `?token=`；
/// 原生客户端（curl / websocat / tokio-tungstenite）可以用任一方式。
fn extract_ws_token(headers: &HeaderMap, query: &HashMap<String, String>) -> Option<String> {
    if let Some(t) = query.get("token") {
        return Some(t.clone());
    }
    headers
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer ").map(|s| s.trim().to_string()))
}

/// 校验 WebSocket 握手的 token；返回 `Ok(Some(rec))` 表示 scope 模式 + 通过；
/// `Ok(None)` 表示 legacy 模式（未配 KeyStore），所有事件无条件放行。
///
/// - `key_store.is_configured() == false` → 无条件放行（legacy 模式）
/// - 配置了 KeyStore：必须有 token，且 key 有 `qot:read` scope（最低门槛，
///   实际收哪些事件由后续 push filter 按 scope 决定）
fn authenticate_ws(
    key_store: &KeyStore,
    headers: &HeaderMap,
    query: &HashMap<String, String>,
) -> Result<Option<Arc<KeyRecord>>, (StatusCode, &'static str)> {
    if !key_store.is_configured() {
        return Ok(None);
    }

    let Some(token) = extract_ws_token(headers, query) else {
        futu_auth::audit::reject(
            "ws",
            "/ws",
            "<missing>",
            "missing token (query or Authorization)",
        );
        return Err((StatusCode::UNAUTHORIZED, "missing api key"));
    };

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

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

    if !rec.scopes.contains(&Scope::QotRead) {
        // v1.4.102 BUG-011 fix (P2): 不再泄露 scope 给请求方,
        // 仅写本地 audit log. 与 REST / gRPC 同步.
        futu_auth::audit::reject("ws", "/ws", &rec.id, "missing qot:read scope");
        return Err((StatusCode::FORBIDDEN, "forbidden"));
    }

    futu_auth::audit::allow("ws", "/ws", &rec.id, Some("qot:read"));
    Ok(Some(rec))
}

/// WebSocket 升级处理
pub async fn ws_handler(
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    Query(query): Query<HashMap<String, String>>,
    State(state): State<RestState>,
) -> impl IntoResponse {
    let rec = match authenticate_ws(&state.key_store, &headers, &query) {
        Ok(rec) => rec,
        Err((code, msg)) => return (code, msg).into_response(),
    };
    // legacy（rec=None）时给个"全 scope"快照让 filter 全放行；scope 模式用 rec.scopes
    let scopes: HashSet<Scope> = match &rec {
        Some(r) => r.scopes.clone(),
        None => all_scopes(),
    };
    let key_id = rec.as_ref().map(|r| r.id.clone());
    // v1.4.102 codex 47 F1 / 48 F1 (P1): WS 必须独立 enforce key.allowed_acc_ids
    // (per-key acc 白名单), 不能依赖 sub-acc-push 是否调过. 之前未调 sub-acc-push
    // 时 fall-back 全 push, key 被限到 acc A 仍能收 acc B 的 trade push.
    let allowed_acc_ids = rec.as_ref().and_then(|r| r.allowed_acc_ids.clone());
    // v1.4.105 D3 (Phase 4) T-B1: per-key allowed_markets 硬限额 (大写字符
    // 串 set, e.g. {"HK","US"}). `None` / 空 set = 无限制. WS Layer 3
    // (TradePushFilter) 用此 set 过滤 trade event 的 trd_market.
    let allowed_markets = rec.as_ref().and_then(|r| r.allowed_markets.clone());
    let broadcaster = Arc::clone(&state.ws_broadcaster);
    // v1.4.102 codex 46 F2 (P1): pass per-key acc subscription state into
    // WS connection so trade push delivery can filter by sub-acc-push registrations.
    let rest_acc_subs = Arc::clone(&state.rest_acc_subscriptions);
    // v1.4.105 D4 (Phase 1): pass shared FilterRegistry into WS connection
    // so push event filter (TradePushFilter) goes through unified registry.
    let filter_registry = Arc::clone(&state.filter_registry);
    let ctx = WsConnectionContext {
        broadcaster,
        scopes,
        key_id,
        allowed_acc_ids,
        allowed_markets,
        rest_acc_subscriptions: rest_acc_subs,
        filter_registry,
    };
    ws.on_upgrade(move |socket| handle_ws_connection(socket, ctx))
        .into_response()
}

/// 全 scope 集合（legacy 模式用）
fn all_scopes() -> HashSet<Scope> {
    [
        Scope::QotRead,
        Scope::AccRead,
        Scope::TradeSimulate,
        Scope::TradeReal,
    ]
    .into_iter()
    .collect()
}

/// 处理单个 WebSocket 连接
///
/// `scopes` 是该连接 key 持有的 scope 集合，用于按 `WsPushScope::required_scope()`
/// 过滤推送事件。例如只有 `qot:read` 的 key 不会收到 `trade` 类推送。
// v1.4.102 codex 47 F1 / 48 F1 (P1): per-key allowed_acc_ids 硬限额.
// `None` = 该 key 无 acc 限制 (默认全开); `Some(set)` = 仅这些 acc 可见.
struct WsConnectionContext {
    broadcaster: Arc<WsBroadcaster>,
    scopes: HashSet<Scope>,
    key_id: Option<String>,
    allowed_acc_ids: Option<HashSet<u64>>,
    // v1.4.105 D3 (Phase 4) T-B1: caller key 的 allowed_markets 硬限额, 用于
    // Layer 3 (TradePushFilter) 过滤. None / 空 set = 无限制.
    allowed_markets: Option<HashSet<String>>,
    rest_acc_subscriptions: Arc<RwLock<HashMap<String, HashSet<u64>>>>,
    // v1.4.105 D4 (Phase 1): 共享 FilterRegistry 实例 — push event 过滤走
    // 同一 registry 与 4 surface (REST body filter / WS push) 一致.
    filter_registry: Arc<futu_auth_pipeline::FilterRegistry>,
}

async fn handle_ws_connection(socket: WebSocket, ctx: WsConnectionContext) {
    let WsConnectionContext {
        broadcaster,
        scopes,
        key_id,
        allowed_acc_ids,
        allowed_markets,
        rest_acc_subscriptions,
        filter_registry,
    } = ctx;

    let (mut ws_tx, mut ws_rx) = socket.split();
    let mut push_rx = broadcaster.subscribe();

    tracing::info!(
        key_id = ?key_id,
        scopes = ?scopes,
        "WebSocket push client connected"
    );

    // v1.4.106 codex 1125 F6 [P2]: REST WS notify subscription state.
    //
    // 对齐 C++ raw TCP `IsConnSubRecvNotify` (APIServer_Qot_PriceReminder.cpp:730-735):
    // broadcast notify 类 push (e.g. price reminder) 必须 client 显式 sub 才下发.
    //
    // **Breaking change vs v1.4.105**: v1.4.105 之前 REST `/ws` 默认收所有
    // broadcast notify; v1.4.106 起需 client 发 `{"action":"subscribe-notify"}`
    // text message 才能继续收. 老 client 如果依赖 price reminder push 必须升级.
    //
    // Default false 对齐 raw TCP 默认 unsub 状态.
    let notify_subscribed = Arc::new(std::sync::atomic::AtomicBool::new(false));
    let notify_subscribed_for_send = Arc::clone(&notify_subscribed);
    let notify_subscribed_for_recv = Arc::clone(&notify_subscribed);

    // 推送任务：从 broadcast channel 读取事件 → 按 scope 过滤 → 发送给客户端
    let send_scopes = scopes.clone();
    let send_key_id_str = key_id.clone().unwrap_or_else(|| "<none>".to_string());
    let send_key_id_for_filter = key_id.clone();
    let rest_subs_for_filter = Arc::clone(&rest_acc_subscriptions);
    let send_task = tokio::spawn(async move {
        while let Ok(event) = push_rx.recv().await {
            // 按 client scope 过滤：key 没这个 scope 就不发
            if !send_scopes.contains(&event.required_scope.required_scope()) {
                // 记一次"被挡住的推送"，供 Prometheus `/metrics` 观察
                futu_auth::metrics::bump_ws_filtered(&event.event_type, &send_key_id_str);
                continue;
            }
            // v1.4.106 codex 1125 F6 [P2]: notify subscribe gate.
            // 对齐 C++ raw TCP `IsConnSubRecvNotify` (broadcast push 必须显式 sub).
            if matches!(event.required_scope, WsPushScope::Notify)
                && !notify_subscribed_for_send.load(std::sync::atomic::Ordering::Relaxed)
            {
                futu_auth::metrics::bump_ws_filtered("notify_unsub", &send_key_id_str);
                continue;
            }
            // v1.4.102 codex 47 F1 / 48 F1 (P1): trade push 进 acc-id 过滤,
            // 两层独立 enforce:
            // 1. **key.allowed_acc_ids 硬限额** (Some(set)): event.acc_id 不在
            //    set → drop. 与 sub-acc-push 是否调过无关 (老 key 没 sub 也强限).
            // 2. **REST sub-acc-push state map** (sub_state): 仅当 key 已调过
            //    sub-acc-push 才生效. entry 存在但不含 acc_id → drop. 未调过 →
            //    pass (向后兼容老 client).
            //
            // codex 48 F2 P1 fix: REST sub state empty entry (Some(set) 但 set
            // 空) 也算 "已 unsub all" tombstone, 不允许 fall back 到全 push.
            //
            // v1.4.103 codex F5.11 (P2) round 5: 抽 logic 到
            // `should_drop_trade_event_for_caller` pure fn 让单测可验证.
            //
            // v1.4.105 D4 (Phase 1): 改走 `FilterRegistry::should_drop_event`
            // 让 4 surface (REST `/ws` 现接 + 后续 gRPC subscribe_push 等)
            // 共用同一 registry instance. 防 sibling-route bypass —
            // 任何人加新 push event filter 只在 registry 注册一次, 不需
            // 改各 surface inline. `should_drop_trade_event_for_caller`
            // pure fn 仍保留作 unit test 直接验证 logic, 但 production 走 registry.
            if matches!(event.required_scope, WsPushScope::Trade)
                && let Some(event_acc) = event.acc_id
            {
                let sub_state_owned: Option<HashSet<u64>> =
                    send_key_id_for_filter.as_ref().and_then(|kid| {
                        crate::adapter::with_rest_acc_subscriptions_read(
                            &rest_subs_for_filter,
                            |subs| subs.get(kid).cloned(),
                        )
                    });
                let ctx = futu_auth_pipeline::PushEventCtx {
                    event_type: &event.event_type,
                    event_acc: Some(event_acc),
                    allowed_acc_ids: allowed_acc_ids.as_ref(),
                    sub_state: sub_state_owned.as_ref(),
                    // v1.4.105 D3 (Phase 4) T-B1: 真接 trd_market —
                    // PushDispatcher 端一次 decode 后透传到 WsPushEvent.trd_market
                    // (None = 老路径 / decode 失败 / market 未知, 不 trigger
                    // Layer 3 drop).
                    event_trd_market: event.trd_market.as_deref(),
                    allowed_markets: allowed_markets.as_ref(),
                };
                if filter_registry.should_drop_event(&ctx) {
                    // v1.4.105 F5.2 fix (codex review C4 4th): 4 surface 统一
                    // metric label "trade_market" 跟 gRPC + raw TCP WS + MCP 一致,
                    // 不再用 event.event_type (= "trade") 让跨 surface jq aggregate
                    // 一致.
                    futu_auth::metrics::bump_ws_filtered("trade_market", &send_key_id_str);
                    continue;
                }
            }
            let json = match serde_json::to_string(&event) {
                Ok(j) => j,
                Err(_) => continue,
            };
            if ws_tx.send(Message::Text(json.into())).await.is_err() {
                break; // 客户端断开
            }
        }
    });

    // 接收任务：处理客户端消息（ping/pong/close + v1.4.106 codex 1125 F6 subscribe-notify）
    let recv_task = tokio::spawn(async move {
        while let Some(msg) = ws_rx.next().await {
            match msg {
                Ok(Message::Close(_)) | Err(_) => break,
                Ok(Message::Ping(data)) => {
                    // axum 自动回复 pong，不需要手动处理
                    let _ = data;
                }
                // v1.4.106 codex 1125 F6 [P2]: 处理 client 发的 JSON control message.
                // 支持 `{"action":"subscribe-notify"}` / `{"action":"unsubscribe-notify"}`.
                // 对齐 C++ raw TCP IsConnSubRecvNotify 的 sub/unsub 接口.
                Ok(Message::Text(text)) => {
                    if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text)
                        && let Some(action) = val.get("action").and_then(|v| v.as_str())
                    {
                        match action {
                            "subscribe-notify" => {
                                notify_subscribed_for_recv
                                    .store(true, std::sync::atomic::Ordering::Relaxed);
                                tracing::info!("WS client subscribed notify push");
                            }
                            "unsubscribe-notify" => {
                                notify_subscribed_for_recv
                                    .store(false, std::sync::atomic::Ordering::Relaxed);
                                tracing::info!("WS client unsubscribed notify push");
                            }
                            other => {
                                tracing::debug!(action = %other, "WS client unknown action");
                            }
                        }
                    }
                }
                _ => {} // 忽略其他消息
            }
        }
    });

    // 任一任务结束则关闭连接
    tokio::select! {
        _ = send_task => {}
        _ = recv_task => {}
    }

    tracing::info!("WebSocket push client disconnected");
}

/// v1.4.103 codex F5.11 (P2) round 5: pure-fn 提取 trade event 过滤决策让单测
/// 可以验证. 不直接 hit handle_ws_connection 异步通路 (需 WebSocket infra),
/// 但 logic 与该函数 inline 顺序 1:1 对齐.
///
/// **行为**:
/// - event_acc=None (非 trade event 没 acc_id): 不 drop
/// - Layer 1: `allowed_acc_ids` 非空 + event_acc ∉ allowed → drop
/// - Layer 2: `sub_state` 含 caller key 的 entry (即使空 tombstone) + event_acc ∉
///   set → drop. **空 set = unsub-all tombstone, 全 drop 一切 trade**
/// - Layer 2 missing entry (caller 从未调过 sub-acc-push) → 不 drop (向后兼容).
///
/// 返 `true` 表示 drop, `false` 表示 deliver.
///
/// v1.4.105 D4 (Phase 1): production 路径已切到 `FilterRegistry::should_drop_event`
/// (走 `TradePushFilter` impl, logic 等价). 这个 pure fn 现仅作 unit test target
/// 直接验证 logic, 不再 production 调用. 保留 + `#[cfg(test)]` 让 dead_code 不警告.
#[cfg(test)]
pub(crate) fn should_drop_trade_event_for_caller(
    allowed_acc_ids: Option<&HashSet<u64>>,
    sub_state: Option<&HashSet<u64>>,
    event_acc: u64,
) -> bool {
    // Layer 1: hard allowed_acc_ids whitelist (caller key 限制).
    if let Some(allowed) = allowed_acc_ids
        && !allowed.is_empty()
        && !allowed.contains(&event_acc)
    {
        return true;
    }
    // Layer 2: per-key REST sub state (opt-in via /api/sub-acc-push).
    // - entry 存在 (sub-acc-push 调过): 必须 acc_id ∈ set (空 set =
    //   unsub all tombstone → drop 一切)
    // - entry 不存在 (从未调 sub-acc-push): 不 drop (向后兼容老 client)
    if let Some(set) = sub_state
        && !set.contains(&event_acc)
    {
        return true;
    }
    false
}

#[cfg(test)]
mod tests;