futu_mcp/state/

push_subscribers.rs

use std::collections::HashSet;
use std::time::Instant;

use rmcp::{RoleServer, service::Peer};

use super::ServerState;
use super::push_filter::subscriber_visible_to_caller;

/// v1.4.38 Phase 5: 订阅了 push 通知的 MCP 客户端 session 记录。
///
/// 每个 session 用 `futu_sub_acc_push` 注册时登记一条。daemon push 事件来到
/// MCP server 时，按 `acc_ids` 过滤后用 `peer.notify_logging_message()` 推回。
///
/// v1.4.38 100%：`acc_ids` 过滤已生效（state.rs drain loop 实装）。
/// caller key ownership / scope 快照在注册时解析并保存，后续 push 分发不再
/// 重新读取 bearer 明文。
#[derive(Clone)]
pub(super) struct PushSubscriber {
    /// rmcp 对 MCP session 的抽象。clone 便宜（内部 Arc）。
    pub peer: Peer<RoleServer>,
    /// 该 session 关心的 acc_id 列表。空集合表示"不过滤"（接收所有 acc 的 push）。
    pub acc_ids: HashSet<u64>,
    /// v1.4.39 per-key acc_id 白名单**注册时快照**（非 live-reload）。
    ///
    /// Some(set) + non-empty → push 的 acc_id 必须在 set 里才推。
    /// Some(empty) / None → 不做 key 级过滤（兼容无 allowed_acc_ids 约束的 key
    /// 或 stdio / legacy 模式）。
    ///
    /// 快照语义：注册后 SIGHUP 重载 keys.json 修改 allowed_acc_ids 不会立即
    /// 反映到已注册订阅者。用户需重新 `futu_sub_acc_push` 才应用新 scope。
    /// 这是 defense-in-depth 层（主 auth 在 tool 调用时 guard.rs），可接受。
    pub allowed_acc_ids_snapshot: Option<HashSet<u64>>,
    /// v1.4.105 T-C2: per-key `allowed_markets` 注册时快照, Layer 3 trade push gate.
    ///
    /// `Some(set)` + non-empty → push event 的 `trd_market` 必须 ∈ set 才推 (按
    /// `Trd_Common.TrdMarket` int → 字符串映射, 与 `keys.json::allowed_markets`
    /// 配置字符串一致, e.g. "HK"/"US"/"FUTURES").
    /// `None` / `Some(empty)` → 不做 market gate (兼容 stdio / legacy / 未配
    /// allowed_markets 的 key).
    ///
    /// 与 `allowed_acc_ids_snapshot` 同样**注册时快照**, SIGHUP 重载不影响
    /// 已注册订阅者. 用户需重新 `futu_sub_acc_push` 才应用新 scope.
    /// 配套 main auth (guard.rs / require_acc_read_with_acc_id) 仍在 tool 调用
    /// 时 enforce, 这是 defense-in-depth 层 (push 走 server-initiated channel
    /// 绕过 tool 调用 → 必须独立 enforce).
    pub allowed_markets_snapshot: Option<HashSet<String>>,
    /// 注册时间（用于 session 硬上限清理，4h 默认 TTL）
    pub registered_at: Instant,
    /// v1.4.103 (codex 50 F6 / 53 F4 — B8): owner key id (KeyRecord.id).
    /// 注册时填的 caller key id (HTTP Bearer 或 startup key); 用于 unsub
    /// ownership check — 任何 caller 拿到 session_id 后想 unsub 必须 key id
    /// 匹配 owner_key_id (admin scope 例外).
    ///
    /// None = legacy / stdio 模式无 key (ownership 退化为 "anyone can unsub",
    /// 与本来 v1.4.102 行为一致).
    pub owner_key_id: Option<String>,
}

impl ServerState {
    /// v1.4.38 Phase 5: 注册当前 session 接收指定 acc_id 的 push。返回 session
    /// UUID（调用方存着，后续可 unregister）。
    ///
    /// v1.4.38: 已 wire 到 `futu_sub_acc_push` tool。tool 被调用时拿到
    /// `RequestContext.peer`，`acc_ids` 从工具 args 解析，注册完成后
    /// state.rs 的 push drain loop 会按 acc_ids filter 转 notify 给该 peer。
    /// v1.4.103 (codex 50 F5 / 53 F2 / 58 F3 — B7) + (codex 50 F6 / 53 F4 — B8):
    /// 注册当前 session 接收指定 acc_id 的 push。
    ///
    /// `owner_key_id_override` 由 caller 传入 (例如从 HTTP Bearer 解析得到 key id);
    /// 若 None → fall back 到 bearer_token 解析 → fall back 到 startup `authed_key.id`.
    /// `allowed_acc_ids_snapshot` 同样 fall back 链: bearer → startup.
    pub async fn register_push_subscriber_with_owner(
        &self,
        peer: Peer<RoleServer>,
        acc_ids: HashSet<u64>,
        bearer_token: Option<String>,
        owner_key_id_override: Option<String>,
    ) -> String {
        let session_id = format!("sub-{}", rand::random::<u64>());

        // 解析 caller-specific KeyRecord (HTTP Bearer 优先, fallback startup key).
        // 用于 (a) allowed_acc_ids_snapshot (B7) (b) owner_key_id (B8).
        //
        // v1.4.106 codex 0608 F2 (P1): startup fallback 路径用
        // `get_by_id_for_current_machine` 替代裸 `get_by_id`, 与 verify (Bearer 路径
        // 自带 machine 校验) 行为对称, 让 SIGHUP 收紧 allowed_machines 后能立即拒绝.
        let bearer_key_rec = bearer_token
            .as_deref()
            .filter(|pt| !pt.is_empty())
            .and_then(|pt| self.key_store.verify(pt));
        let startup_key_rec = self
            .authed_key
            .as_ref()
            .and_then(|k| self.key_store.get_by_id_for_current_machine(&k.id));
        let effective_key_rec = bearer_key_rec.as_ref().or(startup_key_rec.as_ref());

        // v1.4.103 (B7): allowed_acc_ids_snapshot 优先 HTTP Bearer 解析,
        // 否则 fall back startup key.allowed_acc_ids; 都无 → None (无限制).
        let allowed_acc_ids_snapshot =
            effective_key_rec.and_then(|rec| rec.allowed_acc_ids.clone());

        // v1.4.105 T-C2: allowed_markets_snapshot 同链 — 与 acc_ids_snapshot
        // 同样从 effective_key_rec (Bearer / startup) 取 allowed_markets.
        // 用于 push event Layer 3 market gate (state.rs::subscriber_should_receive_with_market).
        let allowed_markets_snapshot =
            effective_key_rec.and_then(|rec| rec.allowed_markets.clone());

        // v1.4.103 (B8): owner_key_id 优先 caller 显式传入, 否则 effective_key_rec.id.
        let owner_key_id =
            owner_key_id_override.or_else(|| effective_key_rec.map(|rec| rec.id.clone()));

        let subscriber = PushSubscriber {
            peer,
            acc_ids,
            allowed_acc_ids_snapshot,
            allowed_markets_snapshot,
            registered_at: Instant::now(),
            owner_key_id,
        };
        self.push_subscribers
            .lock()
            .await
            .insert(session_id.clone(), subscriber);
        session_id
    }

    /// v1.4.103 (codex 50 F6 / 53 F4 — B8): unsub session ownership check.
    ///
    /// 行为:
    /// - 无 caller_key_id (legacy / stdio): 退化为旧行为 (任何 caller 可 unsub).
    /// - 有 caller_key_id + subscriber.owner_key_id 匹配: 删除, 返 Ok(true).
    /// - 有 caller_key_id + subscriber.owner_key_id 不匹配: **拒绝**, 返
    ///   Err(reason) — 防其他 caller 拿可见 session_id 强制 unsub.
    /// - session_id 不存在: 返 Ok(false) (idempotent, 不报错).
    /// - subscriber.owner_key_id = None (legacy 注册): 退化为旧行为 — 任何 caller
    ///   可 unsub (向后兼容).
    pub async fn unregister_push_subscriber_with_owner_check(
        &self,
        session_id: &str,
        caller_key_id: Option<&str>,
    ) -> Result<bool, String> {
        let mut subs = self.push_subscribers.lock().await;
        // 不存在 → idempotent Ok(false), 不报错
        let Some(sub) = subs.get(session_id) else {
            return Ok(false);
        };
        // ownership check
        match (caller_key_id, sub.owner_key_id.as_deref()) {
            (None, _) => {
                // 无 caller key: legacy / stdio 模式, 退化旧行为
                subs.remove(session_id);
                Ok(true)
            }
            (Some(caller), None) => {
                // session 注册时无 owner_key_id (legacy): 任何 caller 可 unsub
                tracing::warn!(
                    session_id,
                    caller,
                    "v1.4.103 B8: unsub legacy session (no owner_key_id) — \
                     allowed for backward-compat"
                );
                subs.remove(session_id);
                Ok(true)
            }
            (Some(caller), Some(owner)) if caller == owner => {
                subs.remove(session_id);
                Ok(true)
            }
            (Some(caller), Some(owner)) => {
                // ownership 不匹配: reject. 当前 MCP subscription ownership contract
                // 没有 admin override surface；若要扩展，必须先在 spec / auth
                // pipeline / integration tests 中定义清楚。
                Err(format!(
                    "session_id {session_id:?} owned by key_id {owner:?}, \
                     caller key_id {caller:?} not allowed to unsub"
                ))
            }
        }
    }

    /// 当前活跃订阅数。生产诊断使用 `push_subscribers_summary`，这里只作为
    /// state 单测的轻量断言入口保留。
    #[cfg(test)]
    pub async fn push_subscriber_count(&self) -> usize {
        self.push_subscribers.lock().await.len()
    }

    /// v1.4.58 Phase A: 列出所有 push 订阅 summary（tool diagnostic 用）。
    ///
    /// 返 Vec<(session_id, acc_ids, age_secs)>。
    ///
    /// **MED-NEW-3 修（2nd review）**：加 `caller_allowed_acc_ids` 参数做
    /// scope-mode 多租过滤。当 caller 的 key 有 `allowed_acc_ids` 白名单时，
    /// **只返** subscription 的 `acc_ids` 与 caller 白名单有交集的条目。
    /// 避免 agent A（acc_ids=[100, 200]）通过本 tool 看到 agent B 订阅的
    /// acc_id=[300, 400]。
    ///
    /// `caller_allowed_acc_ids=None` / empty → 不过滤（legacy mode / no-scope key）。
    ///
    /// **rmcp 版本兼容**：rmcp 1.4.0 `Peer<RoleServer>` 不实装 `PartialEq`，
    /// 无法按 peer 身份直接过滤。若未来 rmcp 加 PartialEq，可切到更精确的
    /// per-session-owner 过滤（当前只能靠 acc_id 权限交集近似）。
    pub async fn push_subscribers_summary(
        &self,
        caller_allowed_acc_ids: Option<&HashSet<u64>>,
    ) -> Vec<(String, HashSet<u64>, u64)> {
        let subs = self.push_subscribers.lock().await;
        let now = Instant::now();
        subs.iter()
            .filter(|(_, sub)| subscriber_visible_to_caller(&sub.acc_ids, caller_allowed_acc_ids))
            .map(|(id, sub)| {
                let age = now
                    .checked_duration_since(sub.registered_at)
                    .map(|d| d.as_secs())
                    .unwrap_or(0);
                // LOW-3RD-1（3rd code review）：scope mode 下返回的 acc_ids 要与
                // caller allowed 求交集 — 防止 caller=[100] 看到 sub=[100, 999]
                // 时知道 999 这个 acc 存在。sub.acc_ids 空集（subscribe-all）不做
                // 交集（概念上 caller 看到的是"有个 catch-all subscriber"，不泄漏
                // 具体账户信息）。
                let visible_accs = match caller_allowed_acc_ids {
                    Some(allowed) if !allowed.is_empty() && !sub.acc_ids.is_empty() => {
                        sub.acc_ids.intersection(allowed).copied().collect()
                    }
                    _ => sub.acc_ids.clone(),
                };
                (id.clone(), visible_accs, age)
            })
            .collect()
    }
}