1use std::collections::{HashMap, HashSet};
8use std::net::SocketAddr;
9use std::sync::{Arc, RwLock};
10
11use axum::extract::connect_info::ConnectInfo;
12use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
13use axum::extract::{Query, State};
14use axum::http::{HeaderMap, StatusCode};
15use axum::response::IntoResponse;
16use chrono::Utc;
17use futures::{SinkExt, StreamExt};
18use tokio::sync::broadcast;
19
20use futu_auth::{KeyRecord, KeyStore, Scope};
21use futu_server::push::ExternalPushSink;
22
23use crate::adapter::RestState;
24
25pub const REST_WS_MAX_CONTROL_MESSAGE_SIZE_BYTES: usize = 64 * 1024;
29
30#[derive(Clone, Debug, serde::Serialize)]
32pub struct WsPushEvent {
33 #[serde(rename = "type")]
35 pub event_type: String,
36 #[serde(skip)]
38 pub required_scope: WsPushScope,
39 pub proto_id: u32,
41 #[serde(skip_serializing_if = "Option::is_none")]
43 pub sec_key: Option<String>,
44 #[serde(skip_serializing_if = "Option::is_none")]
46 pub sub_type: Option<i32>,
47 #[serde(skip_serializing_if = "Option::is_none")]
51 pub rehab_type: Option<i32>,
52 #[serde(skip_serializing_if = "Option::is_none")]
54 pub acc_id: Option<u64>,
55 pub body_b64: String,
57 #[serde(skip_serializing_if = "Option::is_none")]
66 pub trd_market: Option<String>,
67}
68
69#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
75#[non_exhaustive]
76pub enum WsPushScope {
77 #[default]
80 Quote,
81 Notify,
83 Trade,
85}
86
87impl WsPushScope {
88 pub fn required_scope(&self) -> Scope {
90 match self {
91 WsPushScope::Quote => Scope::QotRead,
92 WsPushScope::Notify => Scope::QotRead,
93 WsPushScope::Trade => Scope::AccRead,
94 }
95 }
96}
97
98#[derive(Clone)]
104pub struct WsBroadcaster {
105 tx: broadcast::Sender<WsPushEvent>,
106}
107
108impl WsBroadcaster {
109 pub fn new(capacity: usize) -> Self {
110 let (tx, _) = broadcast::channel(capacity);
111 Self { tx }
112 }
113
114 fn has_receivers(&self) -> bool {
115 self.tx.receiver_count() > 0
116 }
117
118 pub fn send(&self, event: WsPushEvent) {
120 if !self.has_receivers() {
121 return;
122 }
123 let proto_id = event.proto_id;
124 let event_type = event.event_type.clone();
125 if self.tx.send(event).is_err() {
126 tracing::debug!(
127 proto_id,
128 event_type,
129 receiver_count = self.tx.receiver_count(),
130 "rest ws broadcast send skipped"
131 );
132 }
133 }
134
135 pub fn subscribe(&self) -> broadcast::Receiver<WsPushEvent> {
137 self.tx.subscribe()
138 }
139
140 fn encode_body(body: &[u8]) -> String {
141 use base64::Engine;
142 base64::engine::general_purpose::STANDARD.encode(body)
143 }
144
145 pub fn push_quote(
153 &self,
154 sec_key: &str,
155 sub_type: i32,
156 rehab_type: i32,
157 proto_id: u32,
158 body: &[u8],
159 ) {
160 if !self.has_receivers() {
161 return;
162 }
163 self.send(WsPushEvent {
164 event_type: "quote".to_string(),
165 required_scope: WsPushScope::Quote,
166 proto_id,
167 sec_key: Some(sec_key.to_string()),
168 sub_type: Some(sub_type),
169 rehab_type: Some(rehab_type),
170 acc_id: None,
171 body_b64: Self::encode_body(body),
172 trd_market: None,
173 });
174 }
175
176 pub fn push_broadcast(&self, proto_id: u32, body: &[u8]) {
178 if !self.has_receivers() {
179 return;
180 }
181 self.send(WsPushEvent {
182 event_type: "notify".to_string(),
183 required_scope: WsPushScope::Notify,
184 proto_id,
185 sec_key: None,
186 sub_type: None,
187 rehab_type: None,
188 acc_id: None,
189 body_b64: Self::encode_body(body),
190 trd_market: None,
191 });
192 }
193
194 pub fn push_trade(&self, acc_id: u64, proto_id: u32, body: &[u8], trd_market: Option<&str>) {
200 if !self.has_receivers() {
201 return;
202 }
203 self.send(WsPushEvent {
204 event_type: "trade".to_string(),
205 required_scope: WsPushScope::Trade,
206 proto_id,
207 sec_key: None,
208 sub_type: None,
209 rehab_type: None,
210 acc_id: Some(acc_id),
211 body_b64: Self::encode_body(body),
212 trd_market: trd_market.map(|s| s.to_string()),
213 });
214 }
215}
216
217impl ExternalPushSink for WsBroadcaster {
219 fn on_quote_push(
220 &self,
221 sec_key: &str,
222 sub_type: i32,
223 rehab_type: i32,
224 proto_id: u32,
225 body: &[u8],
226 ) {
227 self.push_quote(sec_key, sub_type, rehab_type, proto_id, body);
228 }
229
230 fn on_broadcast_push(&self, proto_id: u32, body: &[u8]) {
231 self.push_broadcast(proto_id, body);
232 }
233
234 fn on_trade_push(&self, acc_id: u64, proto_id: u32, body: &[u8], trd_market: Option<&str>) {
235 self.push_trade(acc_id, proto_id, body, trd_market);
236 }
237}
238
239fn extract_ws_token(headers: &HeaderMap, query: &HashMap<String, String>) -> Option<String> {
244 if let Some(t) = query.get("token") {
245 return Some(t.clone());
246 }
247 headers
248 .get("authorization")
249 .and_then(|v| v.to_str().ok())
250 .and_then(|v| futu_auth_pipeline::parse_bearer_scheme(v).map(|s| s.to_string()))
251}
252
253fn authenticate_ws(
260 key_store: &KeyStore,
261 headers: &HeaderMap,
262 query: &HashMap<String, String>,
263) -> Result<Option<Arc<KeyRecord>>, (StatusCode, &'static str)> {
264 if !key_store.is_configured() {
265 return Ok(None);
266 }
267
268 let Some(token) = extract_ws_token(headers, query) else {
269 futu_auth::audit::reject(
270 "ws",
271 "/ws",
272 "<missing>",
273 "missing token (query or Authorization)",
274 );
275 return Err((StatusCode::UNAUTHORIZED, "missing api key"));
276 };
277
278 let Some(rec) = key_store.verify(&token) else {
279 futu_auth::audit::reject("ws", "/ws", "<invalid>", "invalid api key");
280 return Err((StatusCode::UNAUTHORIZED, "invalid api key"));
281 };
282
283 if rec.is_expired(Utc::now()) {
284 futu_auth::audit::reject("ws", "/ws", &rec.id, "key expired");
285 return Err((StatusCode::UNAUTHORIZED, "key expired"));
286 }
287
288 if !rec.scopes.contains(&Scope::QotRead) {
289 futu_auth::audit::reject("ws", "/ws", &rec.id, "missing qot:read scope");
292 return Err((StatusCode::FORBIDDEN, "forbidden"));
293 }
294
295 futu_auth::audit::allow("ws", "/ws", &rec.id, Some("qot:read"));
296 Ok(Some(rec))
297}
298
299fn headers_have_valid_websocket_key(headers: &HeaderMap) -> bool {
300 futu_auth::websocket::is_valid_sec_websocket_key(
301 headers
302 .get_all("sec-websocket-key")
303 .iter()
304 .map(|value| value.as_bytes()),
305 )
306}
307
308pub async fn ws_handler(
310 ws: WebSocketUpgrade,
311 ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
312 headers: HeaderMap,
313 Query(query): Query<HashMap<String, String>>,
314 State(state): State<RestState>,
315) -> impl IntoResponse {
316 if !headers_have_valid_websocket_key(&headers) {
317 return (StatusCode::BAD_REQUEST, "invalid websocket handshake").into_response();
318 }
319 let peer_addr_string = peer_addr.to_string();
320 let session_id = headers
321 .get("x-request-id")
322 .or_else(|| headers.get("x-futu-session-id"))
323 .and_then(|v| v.to_str().ok())
324 .map(str::trim)
325 .filter(|v| !v.is_empty());
326 let audit_ctx =
327 futu_auth::audit::AuditContext::new(Some(peer_addr_string.as_str()), session_id);
328 let rec = match futu_auth::audit::with_context(audit_ctx.clone(), || {
329 authenticate_ws(&state.key_store, &headers, &query)
330 }) {
331 Ok(rec) => rec,
332 Err((code, msg)) => return (code, msg).into_response(),
333 };
334 let scopes: HashSet<Scope> = match &rec {
336 Some(r) => r.scopes.clone(),
337 None => all_scopes(),
338 };
339 let key_id = rec.as_ref().map(|r| r.id.clone());
340 let allowed_acc_ids = rec.as_ref().and_then(|r| r.allowed_acc_ids.clone());
344 let allowed_markets = rec.as_ref().and_then(|r| r.allowed_markets.clone());
348 let broadcaster = Arc::clone(&state.ws_broadcaster);
349 let rest_acc_subs = Arc::clone(&state.rest_acc_subscriptions);
352 let filter_registry = Arc::clone(&state.filter_registry);
355 let startup_readiness = state.router.startup_readiness();
356 let ctx = WsConnectionContext {
357 broadcaster,
358 scopes,
359 key_id,
360 allowed_acc_ids,
361 allowed_markets,
362 rest_acc_subscriptions: rest_acc_subs,
363 filter_registry,
364 startup_readiness,
365 };
366 ws.max_message_size(REST_WS_MAX_CONTROL_MESSAGE_SIZE_BYTES)
367 .max_frame_size(REST_WS_MAX_CONTROL_MESSAGE_SIZE_BYTES)
368 .on_upgrade(move |socket| handle_ws_connection(socket, ctx))
369 .into_response()
370}
371
372fn all_scopes() -> HashSet<Scope> {
374 [
375 Scope::QotRead,
376 Scope::AccRead,
377 Scope::TradeSimulate,
378 Scope::TradeReal,
379 ]
380 .into_iter()
381 .collect()
382}
383
384struct WsConnectionContext {
391 broadcaster: Arc<WsBroadcaster>,
392 scopes: HashSet<Scope>,
393 key_id: Option<String>,
394 allowed_acc_ids: Option<HashSet<u64>>,
395 allowed_markets: Option<HashSet<String>>,
398 rest_acc_subscriptions: Arc<RwLock<HashMap<String, HashSet<u64>>>>,
399 filter_registry: Arc<futu_auth_pipeline::FilterRegistry>,
402 startup_readiness: futu_server::identity::StartupReadiness,
403}
404
405async fn handle_ws_connection(socket: WebSocket, ctx: WsConnectionContext) {
406 let WsConnectionContext {
407 broadcaster,
408 scopes,
409 key_id,
410 allowed_acc_ids,
411 allowed_markets,
412 rest_acc_subscriptions,
413 filter_registry,
414 startup_readiness,
415 } = ctx;
416
417 let (mut ws_tx, mut ws_rx) = socket.split();
418 let mut push_rx = broadcaster.subscribe();
419
420 tracing::info!(
421 key_id = ?key_id,
422 scopes = ?scopes,
423 "WebSocket push client connected"
424 );
425
426 let notify_subscribed = Arc::new(std::sync::atomic::AtomicBool::new(false));
437 let notify_subscribed_for_send = Arc::clone(¬ify_subscribed);
438 let notify_subscribed_for_recv = Arc::clone(¬ify_subscribed);
439
440 let send_scopes = scopes.clone();
442 let send_key_id_str = key_id.clone().unwrap_or_else(|| "<none>".to_string());
443 let send_key_id_for_filter = key_id.clone();
444 let rest_subs_for_filter = Arc::clone(&rest_acc_subscriptions);
445 let mut send_task = tokio::spawn(async move {
446 loop {
447 let event = match push_rx.recv().await {
448 Ok(event) => event,
449 Err(broadcast::error::RecvError::Lagged(n)) => {
450 tracing::warn!(
451 skipped = n,
452 "REST WebSocket push client lagged, skipped events"
453 );
454 continue;
455 }
456 Err(broadcast::error::RecvError::Closed) => break,
457 };
458 if !rest_ws_push_ready(&startup_readiness) {
459 futu_auth::metrics::bump_ws_filtered("startup_not_ready", &send_key_id_str);
460 continue;
461 }
462 if !send_scopes.contains(&event.required_scope.required_scope()) {
464 futu_auth::metrics::bump_ws_filtered(&event.event_type, &send_key_id_str);
466 continue;
467 }
468 if matches!(event.required_scope, WsPushScope::Notify)
471 && !notify_subscribed_for_send.load(std::sync::atomic::Ordering::Relaxed)
472 {
473 futu_auth::metrics::bump_ws_filtered("notify_unsub", &send_key_id_str);
474 continue;
475 }
476 if matches!(event.required_scope, WsPushScope::Trade)
497 && let Some(event_acc) = event.acc_id
498 {
499 let sub_state_owned: Option<HashSet<u64>> =
500 send_key_id_for_filter.as_ref().and_then(|kid| {
501 crate::adapter::with_rest_acc_subscriptions_read(
502 &rest_subs_for_filter,
503 |subs| subs.get(kid).cloned(),
504 )
505 });
506 let ctx = futu_auth_pipeline::PushEventCtx {
507 event_type: &event.event_type,
508 event_acc: Some(event_acc),
509 allowed_acc_ids: allowed_acc_ids.as_ref(),
510 sub_state: sub_state_owned.as_ref(),
511 event_trd_market: event.trd_market.as_deref(),
516 allowed_markets: allowed_markets.as_ref(),
517 };
518 if filter_registry.should_drop_event(&ctx) {
519 futu_auth::metrics::bump_ws_filtered("trade_market", &send_key_id_str);
524 continue;
525 }
526 }
527 let json = match serde_json::to_string(&event) {
528 Ok(j) => j,
529 Err(_) => continue,
530 };
531 if ws_tx.send(Message::Text(json.into())).await.is_err() {
532 break; }
534 }
535 });
536
537 let mut recv_task = tokio::spawn(async move {
539 while let Some(msg) = ws_rx.next().await {
540 match msg {
541 Ok(Message::Close(_)) | Err(_) => break,
542 Ok(Message::Ping(_data)) => {
543 }
545 Ok(Message::Text(text)) => {
549 if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text)
550 && let Some(action) = val.get("action").and_then(|v| v.as_str())
551 {
552 match action {
553 "subscribe-notify" => {
554 notify_subscribed_for_recv
555 .store(true, std::sync::atomic::Ordering::Relaxed);
556 tracing::info!("WS client subscribed notify push");
557 }
558 "unsubscribe-notify" => {
559 notify_subscribed_for_recv
560 .store(false, std::sync::atomic::Ordering::Relaxed);
561 tracing::info!("WS client unsubscribed notify push");
562 }
563 other => {
564 tracing::debug!(action = %other, "WS client unknown action");
565 }
566 }
567 }
568 }
569 _ => {} }
571 }
572 });
573
574 tokio::select! {
578 _ = &mut send_task => {
579 recv_task.abort();
580 }
581 _ = &mut recv_task => {
582 send_task.abort();
583 }
584 }
585
586 tracing::info!("WebSocket push client disconnected");
587}
588
589fn rest_ws_push_ready(readiness: &futu_server::identity::StartupReadiness) -> bool {
590 readiness.snapshot().state == futu_server::identity::StartupState::Ready
591}
592
593#[cfg(test)]
594mod tests;