Skip to main content

futu_server/
router.rs

1// 请求路由器:按 ProtoID 分发请求到注册的业务处理器
2
3use std::collections::HashMap;
4use std::future::Future;
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use parking_lot::RwLock;
9
10use crate::conn::IncomingRequest;
11use crate::identity::StartupReadiness;
12use crate::protect::ProtectionManager;
13
14/// 业务处理器 trait
15#[async_trait]
16pub trait RequestHandler: Send + Sync + 'static {
17    /// 处理请求,返回响应 body(protobuf 编码后的字节)
18    /// 返回 None 表示不产生响应(例如异步响应场景)
19    async fn handle(&self, conn_id: u64, request: &IncomingRequest) -> Option<Vec<u8>>;
20}
21
22/// 函数形式的处理器包装
23///
24/// 用来把闭包 `Fn(&IncomingRequest) -> Future<Option<Vec<u8>>>` 变成
25/// [`RequestHandler`] trait 对象,避免每个 handler 都写一个 struct + impl。
26pub struct FnHandler<F>(pub F);
27
28#[async_trait]
29impl<F, Fut> RequestHandler for FnHandler<F>
30where
31    F: Fn(u64, bytes::Bytes) -> Fut + Send + Sync + 'static,
32    Fut: Future<Output = Option<Vec<u8>>> + Send + 'static,
33{
34    async fn handle(&self, conn_id: u64, request: &IncomingRequest) -> Option<Vec<u8>> {
35        (self.0)(conn_id, request.body.clone()).await
36    }
37}
38
39/// 请求路由器
40pub struct RequestRouter {
41    handlers: RwLock<HashMap<u32, Arc<dyn RequestHandler>>>,
42    protection: ProtectionManager,
43    startup_readiness: RwLock<StartupReadiness>,
44}
45
46impl RequestRouter {
47    /// 创建空路由器。使用 [`Self::register`] 挂 handler。
48    pub fn new() -> Self {
49        Self::with_startup_readiness(StartupReadiness::default())
50    }
51
52    pub fn with_startup_readiness(startup_readiness: StartupReadiness) -> Self {
53        Self {
54            handlers: RwLock::new(HashMap::new()),
55            protection: ProtectionManager::new(),
56            startup_readiness: RwLock::new(startup_readiness),
57        }
58    }
59
60    pub fn startup_readiness(&self) -> StartupReadiness {
61        self.startup_readiness.read().clone()
62    }
63
64    pub fn set_startup_readiness(&self, startup_readiness: StartupReadiness) {
65        *self.startup_readiness.write() = startup_readiness;
66    }
67
68    /// 注册业务处理器
69    pub fn register(&self, proto_id: u32, handler: Arc<dyn RequestHandler>) {
70        self.handlers.write().insert(proto_id, handler);
71    }
72
73    /// 分发请求到对应处理器
74    pub async fn dispatch(&self, conn_id: u64, request: &IncomingRequest) -> Option<Vec<u8>> {
75        let startup = self.startup_readiness.read().clone();
76        if !startup.allows_proto(request.proto_id) {
77            return Some(make_error_response(-1, "gateway authentication not ready"));
78        }
79        if request.proto_id == futu_core::proto_id::VERIFICATION
80            && startup.snapshot().state != crate::identity::StartupState::Ready
81            && !request.caller_has_auth_setup_scope
82            && !(request.caller_is_loopback && request.caller_legacy_local_mode)
83        {
84            return Some(make_error_response(
85                -1,
86                "verification pre-login admission denied",
87            ));
88        }
89        if self.protection.check_request_freq_limit(conn_id, request) {
90            tracing::warn!(
91                proto_id = request.proto_id,
92                conn_id = conn_id,
93                "request frequency limit exceeded"
94            );
95            return Some(make_error_response(-1, "request frequency limit exceeded"));
96        }
97
98        let handler = {
99            let handlers = self.handlers.read();
100            handlers.get(&request.proto_id).cloned()
101        };
102
103        match handler {
104            Some(h) => {
105                futu_command_runtime::with_surface_command_admission(
106                    request.proto_id,
107                    h.handle(conn_id, request),
108                )
109                .await
110            }
111            None => {
112                tracing::warn!(
113                    proto_id = request.proto_id,
114                    conn_id = conn_id,
115                    "no handler registered"
116                );
117                // 返回通用错误响应
118                Some(make_error_response(-1, "unknown protocol"))
119            }
120        }
121    }
122}
123
124impl Default for RequestRouter {
125    fn default() -> Self {
126        Self::new()
127    }
128}
129
130/// 构造通用错误响应(使用 InitConnect::Response 格式,与 C++ 兼容)
131fn make_error_response(ret_type: i32, msg: &str) -> Vec<u8> {
132    let resp = futu_proto::init_connect::Response {
133        ret_type,
134        ret_msg: Some(msg.to_string()),
135        err_code: None,
136        s2c: None,
137    };
138    prost::Message::encode_to_vec(&resp)
139}
140
141#[cfg(test)]
142mod tests;