1use 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#[async_trait]
16pub trait RequestHandler: Send + Sync + 'static {
17 async fn handle(&self, conn_id: u64, request: &IncomingRequest) -> Option<Vec<u8>>;
20}
21
22pub 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
39pub struct RequestRouter {
41 handlers: RwLock<HashMap<u32, Arc<dyn RequestHandler>>>,
42 protection: ProtectionManager,
43 startup_readiness: RwLock<StartupReadiness>,
44}
45
46impl RequestRouter {
47 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 pub fn register(&self, proto_id: u32, handler: Arc<dyn RequestHandler>) {
70 self.handlers.write().insert(proto_id, handler);
71 }
72
73 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 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
130fn 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;