1use std::collections::HashSet;
4use std::future::Future;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
7use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
8
9use bytes::Bytes;
10use futures::{Sink, SinkExt, Stream, StreamExt};
11use tokio::net::TcpStream;
12use tokio::sync::{mpsc, watch};
13use tokio_util::codec::Framed;
14
15use futu_auth::Scope;
16use futu_codec::FutuCodec;
17use futu_codec::frame::FutuFrame;
18use futu_codec::header::ProtoFmtType;
19use futu_core::error::FutuError;
20use futu_net::encrypt;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum ConnState {
26 Connected,
28 Initialized,
30 Active,
32 Disconnected,
34}
35
36#[derive(Clone)]
42pub(crate) struct ClientCloseControl {
43 tx: watch::Sender<bool>,
44}
45
46pub(crate) struct ClientIoStart {
52 tx: watch::Sender<bool>,
53}
54
55impl ClientIoStart {
56 pub(crate) fn channel() -> (Self, watch::Receiver<bool>) {
57 let (tx, rx) = watch::channel(false);
58 (Self { tx }, rx)
59 }
60
61 pub(crate) fn start(self) {
62 self.tx.send_replace(true);
63 }
64}
65
66pub(crate) async fn wait_for_client_io_start(mut start_rx: watch::Receiver<bool>) -> bool {
67 if *start_rx.borrow() {
68 return true;
69 }
70 start_rx.changed().await.is_ok() && *start_rx.borrow()
71}
72
73pub(crate) async fn wait_for_connection_removal(
74 connections: &dashmap::DashMap<u64, ClientConn>,
75 conn_id: u64,
76) {
77 let mut interval = tokio::time::interval(Duration::from_millis(25));
81 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
82 loop {
83 interval.tick().await;
84 if !connections.contains_key(&conn_id) {
85 return;
86 }
87 }
88}
89
90pub(crate) enum InitConnectStart {
91 Immediate(PreparedInitConnect),
92 Deferred(tokio::task::JoinHandle<()>),
93}
94
95pub(crate) async fn complete_prepared_init_connect(
96 connections: &dashmap::DashMap<u64, ClientConn>,
97 conn_id: u64,
98 serial_no: u32,
99 prepared: PreparedInitConnect,
100) -> Result<(), FutuError> {
101 let tx = connections
102 .get(&conn_id)
103 .ok_or_else(|| FutuError::Codec("InitConnect connection no longer exists".to_string()))?
104 .tx
105 .clone();
106 let frame = FutuFrame::new(
107 futu_core::proto_id::INIT_CONNECT,
108 serial_no,
109 Bytes::from(prepared.response_body.clone()),
110 );
111 tx.send(frame)
112 .await
113 .map_err(|_| FutuError::Codec("InitConnect response connection closed".to_string()))?;
114 let mut conn = connections.get_mut(&conn_id).ok_or_else(|| {
115 FutuError::Codec("InitConnect connection closed after response".to_string())
116 })?;
117 conn.commit_prepared_init_connect(&prepared);
118 Ok(())
119}
120
121async fn send_init_connect_failure(
122 connections: &dashmap::DashMap<u64, ClientConn>,
123 conn_id: u64,
124 serial_no: u32,
125 body: Vec<u8>,
126) -> Result<(), FutuError> {
127 let tx = connections
128 .get(&conn_id)
129 .ok_or_else(|| FutuError::Codec("InitConnect connection no longer exists".to_string()))?
130 .tx
131 .clone();
132 tx.send(FutuFrame::new(
133 futu_core::proto_id::INIT_CONNECT,
134 serial_no,
135 Bytes::from(body),
136 ))
137 .await
138 .map_err(|_| FutuError::Codec("InitConnect failure response connection closed".to_string()))
139}
140
141#[allow(clippy::too_many_arguments)]
142pub(crate) fn start_init_connect_for_startup(
143 connections: Arc<dashmap::DashMap<u64, ClientConn>>,
144 readiness: crate::identity::StartupReadiness,
145 conn_id: u64,
146 body: &[u8],
147 serial_no: u32,
148 server_ver: i32,
149 keepalive_interval: i32,
150 rsa_private_key: Option<String>,
151) -> Result<InitConnectStart, FutuError> {
152 let decoded = ClientConn::decode_init_connect(body, rsa_private_key.as_deref())?;
153 let identity = readiness.snapshot();
154 if decoded.is_internal_ui()
155 || (identity.state == crate::identity::StartupState::Ready && identity.user_id.is_some())
156 {
157 let conn = connections.get(&conn_id).ok_or_else(|| {
161 FutuError::Codec("InitConnect connection no longer exists".to_string())
162 })?;
163 return conn
164 .prepare_decoded_init_connect_with_identity(
165 decoded,
166 server_ver,
167 identity.user_id.unwrap_or(0),
168 identity.attribution,
169 keepalive_interval,
170 rsa_private_key.as_deref(),
171 )
172 .map(InitConnectStart::Immediate);
173 }
174
175 let pending_guard = readiness.track_pending_init_connect();
176 let task = tokio::spawn(async move {
177 let _pending_guard = pending_guard;
178 match readiness
179 .await_init_identity_or_cancel(wait_for_connection_removal(
180 connections.as_ref(),
181 conn_id,
182 ))
183 .await
184 {
185 Ok(Some(identity)) => {
186 let Some(conn) = connections.get(&conn_id) else {
187 return;
188 };
189 let prepared = conn.prepare_decoded_init_connect_with_identity(
190 decoded,
191 server_ver,
192 identity.user_id.unwrap_or(0),
193 identity.attribution,
194 keepalive_interval,
195 rsa_private_key.as_deref(),
196 );
197 drop(conn);
198 match prepared {
199 Ok(prepared) => {
200 if let Err(error) = complete_prepared_init_connect(
201 connections.as_ref(),
202 conn_id,
203 serial_no,
204 prepared,
205 )
206 .await
207 {
208 tracing::warn!(conn_id, error = %error, "deferred InitConnect completion failed");
209 }
210 }
211 Err(error) => {
212 tracing::warn!(conn_id, error = %error, "deferred InitConnect preparation failed");
213 }
214 }
215 }
216 Ok(None) => {}
217 Err(_) => {
218 match ClientConn::encode_init_connect_failure(
219 "Gateway initialization failed",
220 rsa_private_key.as_deref(),
221 ) {
222 Ok(body) => {
223 if let Err(error) = send_init_connect_failure(
224 connections.as_ref(),
225 conn_id,
226 serial_no,
227 body,
228 )
229 .await
230 {
231 tracing::warn!(conn_id, error = %error, "deferred InitConnect failure response failed");
232 }
233 }
234 Err(error) => {
235 tracing::warn!(conn_id, error = %error, "deferred InitConnect failure encoding failed");
236 }
237 }
238 }
239 }
240 });
241 Ok(InitConnectStart::Deferred(task))
242}
243
244impl ClientCloseControl {
245 pub(crate) fn channel() -> (Self, watch::Receiver<bool>) {
246 let (tx, rx) = watch::channel(false);
247 (Self { tx }, rx)
248 }
249
250 pub(crate) fn request_close(&self) -> bool {
256 self.tx.send_if_modified(|closing| {
257 if *closing {
258 false
259 } else {
260 *closing = true;
261 true
262 }
263 })
264 }
265}
266
267pub(crate) async fn await_until_client_close<T>(
268 close_rx: &mut watch::Receiver<bool>,
269 future: impl Future<Output = T>,
270) -> Option<T> {
271 if *close_rx.borrow() {
272 return None;
273 }
274 tokio::select! {
275 result = future => Some(result),
276 changed = close_rx.changed() => {
277 match changed {
278 Ok(()) if !*close_rx.borrow() => {
279 tracing::warn!("client close control changed without entering closing state");
280 }
281 Ok(()) | Err(_) => {}
282 }
283 None
284 }
285 }
286}
287
288pub(crate) async fn await_client_operation<T>(
289 close_rx: &mut watch::Receiver<bool>,
290 future: impl Future<Output = T>,
291 disconnect_tx: &mpsc::UnboundedSender<DisconnectNotify>,
292 conn_id: u64,
293 close_reason: &'static str,
294) -> Option<T> {
295 let result = await_until_client_close(close_rx, future).await;
296 if result.is_none() {
297 notify_disconnect(disconnect_tx, conn_id, close_reason);
298 }
299 result
300}
301
302async fn run_tcp_send_loop<S>(
303 mut frame_rx: mpsc::Receiver<FutuFrame>,
304 mut sink: S,
305 mut close_rx: watch::Receiver<bool>,
306 disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
307 conn_id: u64,
308) where
309 S: Sink<FutuFrame> + Unpin,
310 S::Error: std::fmt::Display,
311{
312 loop {
313 let Some(frame) = await_client_operation(
314 &mut close_rx,
315 frame_rx.recv(),
316 &disconnect_tx,
317 conn_id,
318 "tcp frame receive cancelled",
319 )
320 .await
321 .flatten() else {
322 break;
323 };
324 match await_client_operation(
325 &mut close_rx,
326 sink.send(frame),
327 &disconnect_tx,
328 conn_id,
329 "tcp socket send cancelled",
330 )
331 .await
332 {
333 Some(Ok(())) => {}
334 Some(Err(error)) => {
335 tracing::warn!(conn_id, error = %error, "send failed");
336 notify_disconnect(&disconnect_tx, conn_id, "tcp send failed");
337 break;
338 }
339 None => break,
340 }
341 }
342}
343
344async fn run_tcp_receive_loop<St, E>(
345 mut stream: St,
346 req_tx: mpsc::Sender<IncomingRequest>,
347 mut close_rx: watch::Receiver<bool>,
348 disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
349 mut shutdown_rx: watch::Receiver<bool>,
350 conn_id: u64,
351) where
352 St: Stream<Item = Result<FutuFrame, E>> + Unpin,
353 E: std::fmt::Display,
354{
355 loop {
356 let result = tokio::select! {
357 changed = shutdown_rx.changed() => {
358 if changed.is_err() || *shutdown_rx.borrow() {
359 tracing::info!(
360 conn_id,
361 "connection receive loop stopped by shutdown signal"
362 );
363 break;
364 }
365 continue;
366 }
367 result = await_client_operation(
368 &mut close_rx,
369 stream.next(),
370 &disconnect_tx,
371 conn_id,
372 "tcp socket read cancelled",
373 ) => {
374 let Some(result) = result else {
375 break;
376 };
377 result
378 },
379 };
380 let Some(result) = result else {
381 break;
382 };
383 match result {
384 Ok(frame) => {
385 let req = IncomingRequest::builder(
386 conn_id,
387 frame.header.proto_id,
388 frame.header.serial_no,
389 frame.header.proto_fmt_type,
390 frame.body,
391 )
392 .build();
393 match await_client_operation(
394 &mut close_rx,
395 req_tx.send(req),
396 &disconnect_tx,
397 conn_id,
398 "tcp request forward cancelled",
399 )
400 .await
401 {
402 Some(Ok(())) => {}
403 Some(Err(_)) | None => break,
404 }
405 }
406 Err(error) => {
407 tracing::warn!(conn_id, error = %error, "recv error");
408 break;
409 }
410 }
411 }
412 tracing::info!(conn_id, "connection closed");
413 notify_disconnect(&disconnect_tx, conn_id, "tcp receive loop closed");
414}
415
416include!("conn/model.rs");
417
418impl ClientConn {
419 pub fn generate_conn_id() -> u64 {
421 let millis = conn_id_epoch_elapsed_or_zero().as_millis() as u64;
422 let rand_part: u32 = rand::random();
423 (millis << 22) | (rand_part as u64 & 0x3FFFFF)
424 }
425
426 pub fn generate_session_generation() -> u64 {
427 static NEXT_GENERATION: AtomicU64 = AtomicU64::new(0);
428 NEXT_GENERATION
429 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
430 .wrapping_add(1)
431 .max(1)
432 }
433
434 pub fn generate_aes_key() -> [u8; 16] {
436 let rand_val: u64 = rand::random();
437 let hex = format!("{rand_val:016X}");
438 let mut key = [0u8; 16];
439 key.copy_from_slice(hex.as_bytes());
440 key
441 }
442
443 pub fn make_frame(&self, proto_id: u32, serial_no: u32, body: Bytes) -> FutuFrame {
452 let body_sha1 = FutuFrame::body_sha1(&body);
453 self.make_frame_with_sha1(proto_id, serial_no, body, body_sha1)
454 }
455
456 pub fn make_frame_with_sha1(
458 &self,
459 proto_id: u32,
460 serial_no: u32,
461 body: Bytes,
462 body_sha1: [u8; 20],
463 ) -> FutuFrame {
464 if self.aes_encrypt_enabled {
465 let encrypted = encrypt::aes_ecb_encrypt(&self.aes_key, &body);
466 FutuFrame::with_sha1(proto_id, serial_no, Bytes::from(encrypted), body_sha1)
467 } else {
468 FutuFrame::with_sha1(proto_id, serial_no, body, body_sha1)
469 }
470 }
471
472 pub fn decrypt_body(&self, body: &[u8]) -> Result<Vec<u8>, FutuError> {
476 if self.aes_encrypt_enabled {
477 encrypt::aes_ecb_decrypt(&self.aes_key, body).map_err(|e| {
478 tracing::warn!(conn_id = self.conn_id, error = %e, "AES decrypt body failed");
479 e
480 })
481 } else {
482 Ok(body.to_vec())
483 }
484 }
485
486 pub fn handle_init_connect(
494 &mut self,
495 body: &[u8],
496 server_ver: i32,
497 login_user_id: u64,
498 keepalive_interval: i32,
499 rsa_private_key: Option<&str>,
500 ) -> Result<Vec<u8>, FutuError> {
501 self.handle_init_connect_with_identity(
502 body,
503 server_ver,
504 login_user_id,
505 None,
506 keepalive_interval,
507 rsa_private_key,
508 )
509 }
510
511 pub fn handle_init_connect_with_identity(
512 &mut self,
513 body: &[u8],
514 server_ver: i32,
515 login_user_id: u64,
516 user_attribution: Option<i32>,
517 keepalive_interval: i32,
518 rsa_private_key: Option<&str>,
519 ) -> Result<Vec<u8>, FutuError> {
520 let decoded = Self::decode_init_connect(body, rsa_private_key)?;
521 self.handle_decoded_init_connect_with_identity(
522 decoded,
523 server_ver,
524 login_user_id,
525 user_attribution,
526 keepalive_interval,
527 rsa_private_key,
528 )
529 }
530
531 pub(crate) fn decode_init_connect(
532 body: &[u8],
533 rsa_private_key: Option<&str>,
534 ) -> Result<DecodedInitConnect, FutuError> {
535 let decrypted_body;
537 let req_body = if let Some(rsa_key) = rsa_private_key {
538 decrypted_body =
539 futu_net::encrypt::rsa_private_decrypt_blocks(rsa_key, body).map_err(|e| {
540 tracing::warn!(error = %e, "RSA decrypt InitConnect C2S failed");
541 e
542 })?;
543 tracing::debug!(
544 encrypted_len = body.len(),
545 decrypted_len = decrypted_body.len(),
546 "RSA decrypted InitConnect C2S"
547 );
548 &decrypted_body[..]
549 } else {
550 body
551 };
552
553 let req: futu_proto::init_connect::Request =
554 prost::Message::decode(req_body).map_err(FutuError::Proto)?;
555 Ok(DecodedInitConnect { request: req })
556 }
557
558 pub(crate) fn handle_decoded_init_connect_with_identity(
559 &mut self,
560 decoded: DecodedInitConnect,
561 server_ver: i32,
562 login_user_id: u64,
563 user_attribution: Option<i32>,
564 keepalive_interval: i32,
565 rsa_private_key: Option<&str>,
566 ) -> Result<Vec<u8>, FutuError> {
567 let prepared = self.prepare_decoded_init_connect_with_identity(
568 decoded,
569 server_ver,
570 login_user_id,
571 user_attribution,
572 keepalive_interval,
573 rsa_private_key,
574 )?;
575 let response_body = prepared.response_body.clone();
576 self.commit_prepared_init_connect(&prepared);
577 Ok(response_body)
578 }
579
580 pub(crate) fn prepare_decoded_init_connect_with_identity(
581 &self,
582 decoded: DecodedInitConnect,
583 server_ver: i32,
584 login_user_id: u64,
585 user_attribution: Option<i32>,
586 keepalive_interval: i32,
587 rsa_private_key: Option<&str>,
588 ) -> Result<PreparedInitConnect, FutuError> {
589 let req = decoded.request;
590
591 let aes_key_str = std::str::from_utf8(&self.aes_key)
592 .map_err(|e| FutuError::Codec(format!("invalid InitConnect conn_aes_key: {e}")))?
593 .to_string();
594
595 let resp = futu_proto::init_connect::Response {
596 ret_type: 0,
597 ret_msg: None,
598 err_code: None,
599 s2c: Some(futu_proto::init_connect::S2c {
600 server_ver,
601 login_user_id,
602 conn_id: self.conn_id,
603 conn_aes_key: aes_key_str,
604 keep_alive_interval: keepalive_interval,
605 aes_cb_civ: None,
606 user_attribution,
607 }),
608 };
609
610 Ok(PreparedInitConnect {
611 response_body: Self::encode_init_connect_response(&resp, rsa_private_key)?,
612 recv_notify: req.c2s.recv_notify.unwrap_or(false),
613 ai_type: req.c2s.ai_type.unwrap_or(0),
614 enable_aes: rsa_private_key.is_some(),
615 })
616 }
617
618 pub(crate) fn commit_prepared_init_connect(&mut self, prepared: &PreparedInitConnect) {
619 self.state = ConnState::Initialized;
620 self.recv_notify = prepared.recv_notify;
621 self.ai_type = prepared.ai_type;
622 self.aes_encrypt_enabled = prepared.enable_aes;
623 if prepared.enable_aes {
624 tracing::debug!(conn_id = self.conn_id, "AES body encryption enabled");
625 }
626 }
627
628 pub(crate) fn encode_init_connect_failure(
629 message: &str,
630 rsa_private_key: Option<&str>,
631 ) -> Result<Vec<u8>, FutuError> {
632 let response = futu_proto::init_connect::Response {
633 ret_type: -1,
634 ret_msg: Some(message.to_string()),
635 err_code: None,
636 s2c: None,
637 };
638 Self::encode_init_connect_response(&response, rsa_private_key)
639 }
640
641 fn encode_init_connect_response(
642 response: &futu_proto::init_connect::Response,
643 rsa_private_key: Option<&str>,
644 ) -> Result<Vec<u8>, FutuError> {
645 let resp_body = prost::Message::encode_to_vec(response);
646
647 if let Some(rsa_key) = rsa_private_key {
649 let encrypted = futu_net::encrypt::rsa_public_encrypt_blocks(rsa_key, &resp_body)
650 .map_err(|e| {
651 tracing::warn!(error = %e, "RSA encrypt InitConnect S2C failed");
652 e
653 })?;
654 tracing::debug!(
655 plaintext_len = resp_body.len(),
656 encrypted_len = encrypted.len(),
657 "RSA encrypted InitConnect S2C"
658 );
659 Ok(encrypted)
660 } else {
661 Ok(resp_body)
662 }
663 }
664
665 pub fn handle_keepalive(&self, body: &[u8]) -> Result<Vec<u8>, FutuError> {
671 self.handle_keepalive_at(body, chrono::Utc::now().timestamp())
672 }
673
674 pub fn handle_keepalive_at(
679 &self,
680 body: &[u8],
681 server_now_ts: i64,
682 ) -> Result<Vec<u8>, FutuError> {
683 let _req: futu_proto::keep_alive::Request =
684 prost::Message::decode(body).map_err(FutuError::Proto)?;
685
686 self.keepalive_count.fetch_add(1, Ordering::Relaxed);
687
688 let resp = futu_proto::keep_alive::Response {
689 ret_type: 0,
690 ret_msg: None,
691 err_code: None,
692 s2c: Some(futu_proto::keep_alive::S2c {
693 time: server_now_ts,
694 }),
695 };
696
697 Ok(prost::Message::encode_to_vec(&resp))
698 }
699}
700
701pub struct DisconnectNotify {
703 pub conn_id: u64,
705}
706
707pub(crate) fn notify_disconnect(
709 disconnect_tx: &mpsc::UnboundedSender<DisconnectNotify>,
710 conn_id: u64,
711 reason: &'static str,
712) {
713 if let Err(e) = disconnect_tx.send(DisconnectNotify { conn_id }) {
714 tracing::warn!(
715 conn_id,
716 reason,
717 error = %e,
718 "disconnect cleanup notification failed"
719 );
720 }
721}
722
723pub async fn run_connection(
728 stream: TcpStream,
729 conn_id: u64,
730 aes_key: [u8; 16],
731 req_tx: mpsc::Sender<IncomingRequest>,
732 disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
733 shutdown_rx: watch::Receiver<bool>,
734) -> mpsc::Sender<FutuFrame> {
735 let (frame_tx, _close_control) = run_connection_with_close_control(
736 stream,
737 conn_id,
738 aes_key,
739 req_tx,
740 disconnect_tx,
741 shutdown_rx,
742 )
743 .await;
744 frame_tx
745}
746
747pub(crate) async fn run_connection_with_close_control(
748 stream: TcpStream,
749 conn_id: u64,
750 aes_key: [u8; 16],
751 req_tx: mpsc::Sender<IncomingRequest>,
752 disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
753 shutdown_rx: watch::Receiver<bool>,
754) -> (mpsc::Sender<FutuFrame>, ClientCloseControl) {
755 let (frame_tx, close_control, io_start) = prepare_connection_with_close_control(
756 stream,
757 conn_id,
758 aes_key,
759 req_tx,
760 disconnect_tx,
761 shutdown_rx,
762 )
763 .await;
764 io_start.start();
765 (frame_tx, close_control)
766}
767
768pub(crate) async fn prepare_connection_with_close_control(
769 stream: TcpStream,
770 conn_id: u64,
771 _aes_key: [u8; 16],
772 req_tx: mpsc::Sender<IncomingRequest>,
773 disconnect_tx: mpsc::UnboundedSender<DisconnectNotify>,
774 shutdown_rx: watch::Receiver<bool>,
775) -> (mpsc::Sender<FutuFrame>, ClientCloseControl, ClientIoStart) {
776 let (frame_tx, frame_rx) = mpsc::channel::<FutuFrame>(256);
777 let (close_control, close_rx) = ClientCloseControl::channel();
778 let (io_start, start_rx) = ClientIoStart::channel();
779
780 let framed = Framed::new(stream, FutuCodec);
781 let (sink, stream) = framed.split();
782
783 let send_disconnect_tx = disconnect_tx.clone();
785 let send_close_rx = close_rx.clone();
786 let send_close_control_guard = close_control.clone();
787 let send_start_rx = start_rx.clone();
788 tokio::spawn(async move {
789 if !wait_for_client_io_start(send_start_rx).await {
790 return;
791 }
792 let _send_close_control_guard = send_close_control_guard;
793 run_tcp_send_loop(frame_rx, sink, send_close_rx, send_disconnect_tx, conn_id).await;
794 });
795
796 tokio::spawn(async move {
798 if !wait_for_client_io_start(start_rx).await {
799 return;
800 }
801 run_tcp_receive_loop(
802 stream,
803 req_tx,
804 close_rx,
805 disconnect_tx,
806 shutdown_rx,
807 conn_id,
808 )
809 .await;
810 });
811
812 (frame_tx, close_control, io_start)
813}
814
815#[cfg(test)]
816mod tests;