futu_mcp/

guard.rs

//! Scope 守卫 + 限额检查 + 审计日志
//!
//! 两种模式：
//! - **scope 模式**（`state.is_scope_mode()`）：必须持有 api-key，且 scope 匹配
//! - **legacy 模式**（没配 keys-file）：读工具全放行；写工具走 `enable_trading` /
//!   `allow_real_trading` 两级开关

use std::sync::Arc;

use chrono::Utc;
use futu_auth::{CheckCtx, KeyRecord, LimitOutcome, Scope};
use sha2::{Digest, Sha256};

use crate::state::ServerState;

/// 从 KeyStore 里按 id 取**当前**的 KeyRecord（对齐 SIGHUP 热重载）。
///
/// 如果 startup 时的 authed_key 已被 remove_key 删掉，返回 None → 调用方拒绝。
/// 否则返回存储中最新版的 KeyRecord（scope / limits / expires_at 全新鲜）。
fn current_authed_key(state: &ServerState) -> Option<Arc<KeyRecord>> {
    let startup = state.authed_key.as_ref()?;
    // legacy 模式下 key_store 是 empty()，get_by_id 返回 None，
    // 但 legacy 模式在 guard 入口就分支出去了，不会走到这里
    state.key_store.get_by_id(&startup.id)
}

/// 工具需要的 scope 类别
///
/// Read 类的 scope（qot:read / acc:read）是静态确定的；Trade 类在运行时根据
/// `env=real/simulate` 派发到 `Scope::TradeReal` / `Scope::TradeSimulate`。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolScope {
    /// 只读工具需要的 scope
    Read(Scope),
    /// 交易写工具（具体派发由 `require_trading` 根据 env 决定）
    Trade,
}

/// 工具名 → 所需 scope 的**中央注册表**。
///
/// 新加工具时必须同步更新这里，否则 `require_tool_scope()` 返回 None，
/// handler 会 fail-closed 地拒绝请求（"unknown MCP tool"）。
/// 配套的 `all_known_tools_have_scopes` 测试也会失败。
///
/// 对应 REST 的 `futu_rest::auth::scope_for_path()`，保持一致的 fail-closed 策略。
pub fn scope_for_tool(tool: &str) -> Option<ToolScope> {
    match tool {
        // ----- qot:read -----
        "futu_ping" | "futu_get_quote" | "futu_get_snapshot" | "futu_get_kline"
        | "futu_get_orderbook" | "futu_get_ticker" | "futu_get_rt" | "futu_get_static"
        | "futu_get_broker" | "futu_list_plates" | "futu_plate_stocks" => {
            Some(ToolScope::Read(Scope::QotRead))
        }

        // ----- acc:read -----
        "futu_list_accounts" | "futu_get_funds" | "futu_get_positions" | "futu_get_orders"
        | "futu_get_deals" => Some(ToolScope::Read(Scope::AccRead)),

        // ----- trade:real / trade:simulate（运行时 env 派发） -----
        "futu_place_order" | "futu_modify_order" | "futu_cancel_order" => Some(ToolScope::Trade),

        // fail-closed
        _ => None,
    }
}

/// 基于注册表的只读 scope 守卫；若工具未登记则 fail-closed 拒绝。
///
/// 仅对 `ToolScope::Read(_)` 生效；Trade 类工具仍然走 `require_trading()`。
pub fn require_tool_scope(state: &ServerState, tool: &'static str) -> GuardOutcome {
    match scope_for_tool(tool) {
        Some(ToolScope::Read(s)) => require_scope(state, tool, s),
        Some(ToolScope::Trade) => {
            // 防御性分支：调用方错把 trade 工具丢到 require_tool_scope 里。
            audit(tool, None, "reject", "internal: trade tool misrouted");
            GuardOutcome::Reject(format!(
                "internal error: {tool} is a trade tool, must use require_trading"
            ))
        }
        None => {
            audit(tool, None, "reject", "unknown MCP tool");
            GuardOutcome::Reject(format!("unknown MCP tool {tool:?}"))
        }
    }
}

/// 守卫结果：Allow 或携带拒绝原因的 JSON
pub enum GuardOutcome {
    Allow,
    Reject(String),
}

impl GuardOutcome {
    pub fn into_err_json(self) -> Option<String> {
        match self {
            GuardOutcome::Allow => None,
            GuardOutcome::Reject(msg) => Some(serde_json::json!({ "error": msg }).to_string()),
        }
    }
}

/// 基础 scope 守卫（用于只读工具）
///
/// 返回 `GuardOutcome::Allow` 表示放行。legacy 模式下只读工具全放行。
pub fn require_scope(state: &ServerState, tool: &'static str, needed: Scope) -> GuardOutcome {
    if !state.is_scope_mode() {
        // legacy 行为：只读工具不检查（旧用户兼容）
        audit(tool, None, "allow", "legacy mode, no keys configured");
        return GuardOutcome::Allow;
    }

    if state.authed_key.is_none() {
        audit(tool, None, "reject", "no API key provided");
        return GuardOutcome::Reject(
            "API key required: set FUTU_MCP_API_KEY to a plaintext key listed in keys.json"
                .to_string(),
        );
    }

    // SIGHUP 热重载后用 id 重新 lookup，拿最新 scope/limits/expiry
    let Some(key) = current_authed_key(state) else {
        let id = state
            .authed_key
            .as_ref()
            .map(|k| k.id.clone())
            .unwrap_or_default();
        audit(
            tool,
            Some(&id),
            "reject",
            "key revoked (not in current keys.json)",
        );
        return GuardOutcome::Reject(format!(
            "API key {id:?} has been revoked (not in current keys.json)"
        ));
    };

    // 过期再查一次（防止启动后过期，或 SIGHUP 后 expires_at 被改小）
    if key.is_expired(Utc::now()) {
        audit(tool, Some(&key.id), "reject", "key expired");
        return GuardOutcome::Reject(format!(
            "API key {:?} has expired (expires_at={:?})",
            key.id, key.expires_at
        ));
    }

    if !key.scopes.contains(&needed) {
        audit(
            tool,
            Some(&key.id),
            "reject",
            &format!("missing scope {}", needed),
        );
        return GuardOutcome::Reject(format!(
            "API key {:?} missing required scope {:?}",
            key.id,
            needed.as_str()
        ));
    }

    audit(tool, Some(&key.id), "allow", "scope ok");
    GuardOutcome::Allow
}

/// 交易写守卫：scope + legacy 兼容 + （可选）限额检查 + （可选）per-call key 覆盖
///
/// `env`：`"real"` / `"simulate"`；`ctx` 为 Some 时跑限额检查（下单路径）。
///
/// `override_key` 为 Some 时，本次调用使用这个 key（`KeyStore::verify` 一次性
/// 拿最新 record）而不是 `state.authed_key`；典型用法：MCP 多租户，让 LLM 客户端
/// 每个 tool call 带自己的 key。验证失败 → reject，不回落。若为 None → 用启动时
/// 捕获的 `state.authed_key`（SIGHUP-aware fresh lookup）。
pub fn require_trading(
    state: &ServerState,
    tool: &'static str,
    env: &str,
    ctx: Option<CheckCtx>,
    override_key: Option<&str>,
) -> GuardOutcome {
    let is_real = crate::handlers::trade_write::is_real_env(env);
    let needed_scope = if is_real {
        Scope::TradeReal
    } else {
        Scope::TradeSimulate
    };

    if !state.is_scope_mode() {
        // legacy：两级开关
        if !state.enable_trading {
            audit(tool, None, "reject", "legacy: --enable-trading off");
            return GuardOutcome::Reject(
                "trading tools are disabled. Start futu-mcp with --enable-trading to enable."
                    .to_string(),
            );
        }
        if is_real && !state.allow_real_trading {
            audit(
                tool,
                None,
                "reject",
                "legacy: real env but --allow-real-trading off",
            );
            return GuardOutcome::Reject(
                "real trading is not allowed. Use env=\"simulate\" or restart futu-mcp with --allow-real-trading."
                    .to_string(),
            );
        }
        // legacy 下 override_key 被忽略（没有 KeyStore 可 verify），
        // 但这种配置本身就是"信任所有调用方"，覆盖不覆盖不影响安全语义
        audit(tool, None, "allow", "legacy trading allowed");
        return GuardOutcome::Allow;
    }

    // scope 模式：先解析用哪把 key
    let key = if let Some(plaintext) = override_key.filter(|p| !p.is_empty()) {
        // per-call override：用传入的 plaintext 实时 verify，不走 startup 快照
        match state.key_store.verify(plaintext) {
            Some(rec) => rec,
            None => {
                audit(tool, None, "reject", "per-call api_key invalid");
                return GuardOutcome::Reject(
                    "per-call api_key is invalid (not in keys.json or expired/bound to wrong machine)"
                        .to_string(),
                );
            }
        }
    } else {
        if state.authed_key.is_none() {
            audit(tool, None, "reject", "no API key");
            return GuardOutcome::Reject(
                "API key required for trading tools (set FUTU_MCP_API_KEY, or pass api_key in the tool call)"
                    .to_string(),
            );
        }
        // 同样走 SIGHUP-aware 的 fresh lookup
        match current_authed_key(state) {
            Some(k) => k,
            None => {
                let id = state
                    .authed_key
                    .as_ref()
                    .map(|k| k.id.clone())
                    .unwrap_or_default();
                audit(tool, Some(&id), "reject", "key revoked");
                return GuardOutcome::Reject(format!("API key {id:?} has been revoked"));
            }
        }
    };

    if key.is_expired(Utc::now()) {
        audit(tool, Some(&key.id), "reject", "key expired");
        return GuardOutcome::Reject(format!("API key {:?} has expired", key.id));
    }

    if !key.scopes.contains(&needed_scope) {
        audit(
            tool,
            Some(&key.id),
            "reject",
            &format!("missing scope {}", needed_scope),
        );
        return GuardOutcome::Reject(format!(
            "API key {:?} missing scope {:?}",
            key.id,
            needed_scope.as_str()
        ));
    }

    // 限额检查（仅在提供了 ctx 时执行）
    if let Some(ctx) = ctx {
        match state
            .counters
            .check_and_commit(&key.id, &key.limits(), &ctx, Utc::now())
        {
            LimitOutcome::Allow => {
                audit(tool, Some(&key.id), "allow", "scope + limits ok");
            }
            LimitOutcome::Reject(reason) => {
                audit(tool, Some(&key.id), "reject", &format!("limit: {reason}"));
                return GuardOutcome::Reject(format!("limit check failed: {reason}"));
            }
        }
    } else {
        audit(tool, Some(&key.id), "allow", "scope ok (no limits ctx)");
    }

    GuardOutcome::Allow
}

/// 审计日志：key_id / tool / result / reason
fn audit(tool: &str, key_id: Option<&str>, result: &str, reason: &str) {
    let key_id = key_id.unwrap_or("<none>");
    if result == "reject" {
        futu_auth::audit::reject("mcp", tool, key_id, reason);
    } else {
        futu_auth::audit::allow("mcp", tool, key_id, Some(reason));
    }
}

/// 计算 args 的短哈希（前 8 hex），用于审计日志（不存原始敏感字段）
pub fn args_short_hash(args: &impl serde::Serialize) -> String {
    let j = match serde_json::to_vec(args) {
        Ok(v) => v,
        Err(_) => return "n/a".into(),
    };
    let h = Sha256::digest(&j);
    hex::encode(&h[..4])
}

/// 交易工具执行完后的审计事件：解析 handler 返回的 JSON，success / failure
/// 写入 audit JSONL。`key_id = None` 时用 "<none>" 占位（legacy 模式）。
pub fn emit_trade_outcome(tool: &'static str, key_id: Option<&str>, args_hash: &str, result: &str) {
    let key_id = key_id.unwrap_or("<none>");
    let (outcome, reason) = match serde_json::from_str::<serde_json::Value>(result) {
        Ok(v) => match v.get("error").and_then(|e| e.as_str()) {
            Some(err) => ("failure", Some(err.to_string())),
            None => ("success", None),
        },
        Err(_) => ("unknown", Some("non-json response".to_string())),
    };
    futu_auth::audit::trade("mcp", tool, key_id, args_hash, outcome, reason.as_deref());
}

#[cfg(test)]
mod tests {
    use super::*;

    /// "备忘锁"：维护者加了新工具但忘了更新 `scope_for_tool`，这个测试会挂。
    /// 对应 REST 的 `all_known_routes_have_scopes`。
    #[test]
    fn all_known_tools_have_scopes() {
        let known: &[&str] = &[
            // qot:read
            "futu_ping",
            "futu_get_quote",
            "futu_get_snapshot",
            "futu_get_kline",
            "futu_get_orderbook",
            "futu_get_ticker",
            "futu_get_rt",
            "futu_get_static",
            "futu_get_broker",
            "futu_list_plates",
            "futu_plate_stocks",
            // acc:read
            "futu_list_accounts",
            "futu_get_funds",
            "futu_get_positions",
            "futu_get_orders",
            "futu_get_deals",
            // trade
            "futu_place_order",
            "futu_modify_order",
            "futu_cancel_order",
        ];
        for t in known {
            assert!(
                scope_for_tool(t).is_some(),
                "tool {t:?} is declared but not mapped in scope_for_tool()"
            );
        }
    }

    #[test]
    fn unknown_tool_fails_closed() {
        assert!(scope_for_tool("futu_transfer_money").is_none());
        assert!(scope_for_tool("").is_none());
    }

    // -------- per-call api_key override --------

    use futu_auth::{KeyRecord, KeyStore};
    use std::path::PathBuf;
    use tempfile::TempDir;

    /// 造一个带两把 key 的 KeyStore：k_sim (TradeSimulate), k_real (TradeReal)。
    /// 返回 (store, k_sim_plaintext, k_real_plaintext, k_sim_id, k_real_id, tempdir guard)
    fn mk_two_key_store() -> (Arc<KeyStore>, String, String, String, String, TempDir) {
        let dir = tempfile::tempdir().unwrap();
        let path: PathBuf = dir.path().join("keys.json");
        let (pt_sim, r_sim) = KeyRecord::generate(
            "bot-sim",
            [Scope::TradeSimulate].into_iter().collect(),
            None,
            None,
            None,
        );
        let (pt_real, r_real) = KeyRecord::generate(
            "bot-real",
            [Scope::TradeReal].into_iter().collect(),
            None,
            None,
            None,
        );
        let sid = r_sim.id.clone();
        let rid = r_real.id.clone();
        futu_auth::store::append_key(&path, r_sim).unwrap();
        futu_auth::store::append_key(&path, r_real).unwrap();
        let store = Arc::new(KeyStore::load(&path).unwrap());
        (store, pt_sim, pt_real, sid, rid, dir)
    }

    fn mk_state(store: Arc<KeyStore>, startup_plaintext: &str) -> ServerState {
        let rec = store.verify(startup_plaintext);
        ServerState::new("dummy:0".to_string())
            .with_key_store(store)
            .with_authed_key(rec)
    }

    #[test]
    fn override_key_none_uses_startup_key() {
        let (store, pt_sim, _pt_real, _, _, _dir) = mk_two_key_store();
        let state = mk_state(store, &pt_sim);
        // k_sim 只有 TradeSimulate；env=simulate → allow
        let r = require_trading(&state, "futu_place_order", "simulate", None, None);
        assert!(matches!(r, GuardOutcome::Allow));
        // env=real → reject（k_sim 无 TradeReal）
        let r = require_trading(&state, "futu_place_order", "real", None, None);
        assert!(matches!(r, GuardOutcome::Reject(_)));
    }

    #[test]
    fn override_key_some_valid_uses_that_key_scope() {
        let (store, pt_sim, pt_real, _, _, _dir) = mk_two_key_store();
        // 启动时用 k_sim（只能 simulate），但这次调用带 k_real 的 plaintext
        // → 这次允许 real 交易
        let state = mk_state(store, &pt_sim);
        let r = require_trading(
            &state,
            "futu_place_order",
            "real",
            None,
            Some(pt_real.as_str()),
        );
        assert!(matches!(r, GuardOutcome::Allow), "override 应让 real 通过");
    }

    #[test]
    fn override_key_invalid_is_rejected_not_fallback() {
        let (store, pt_sim, _pt_real, _, _, _dir) = mk_two_key_store();
        let state = mk_state(store, &pt_sim);
        // 传一个不在 keys.json 里的假 plaintext
        let r = require_trading(
            &state,
            "futu_place_order",
            "simulate",
            None,
            Some("deadbeef"),
        );
        // 必须拒绝，不能回落到 startup key —— 否则攻击者塞任何字串都能"升级"到 startup key
        assert!(matches!(r, GuardOutcome::Reject(_)));
        if let GuardOutcome::Reject(msg) = r {
            assert!(
                msg.contains("per-call api_key is invalid"),
                "错误文案: {msg}"
            );
        }
    }

    #[test]
    fn override_key_empty_string_treated_as_none() {
        // 用户传了 api_key="" → 等同于没传，应该回落到 startup key
        let (store, pt_sim, _pt_real, _, _, _dir) = mk_two_key_store();
        let state = mk_state(store, &pt_sim);
        let r = require_trading(&state, "futu_place_order", "simulate", None, Some(""));
        assert!(matches!(r, GuardOutcome::Allow));
    }

    #[test]
    fn override_key_scope_mismatch_is_rejected() {
        // override 到 k_sim（只有 TradeSimulate）但想 trade real → reject
        let (store, pt_sim, _pt_real, _, _, _dir) = mk_two_key_store();
        let state = mk_state(store, &pt_sim);
        let r = require_trading(
            &state,
            "futu_place_order",
            "real",
            None,
            Some(pt_sim.as_str()),
        );
        assert!(matches!(r, GuardOutcome::Reject(_)));
    }
}