Skip to main content

futu_server/
conn.rs

1// 单连接管理:状态机、帧收发、加密、心跳超时
2
3use 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/// 连接状态
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum ConnState {
26    /// TCP / WebSocket 刚建立,尚未完成 InitConnect 握手
27    Connected,
28    /// 已完成 InitConnect,等待首次业务请求
29    Initialized,
30    /// 正常交互中(业务请求 / KeepAlive / push 均已流转)
31    Active,
32    /// 连接已断开(客户端主动关闭 / 被动超时 / IO 错误)
33    Disconnected,
34}
35
36/// Per-connection cancellation signal shared by push fanout and socket loops.
37///
38/// A full ordinary-push queue transitions this control from open to closing.
39/// TCP/WS send and receive loops observe the same signal, terminate their
40/// pending I/O, and notify the listener's existing idempotent cleanup path.
41#[derive(Clone)]
42pub(crate) struct ClientCloseControl {
43    tx: watch::Sender<bool>,
44}
45
46/// One-shot registration barrier for a prepared client I/O driver.
47///
48/// Socket send/receive tasks are created before the listener owns the
49/// connection state, but neither task may touch the socket until the listener
50/// inserts [`ClientConn`] and registers its close control.
51pub(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    // Pending InitConnect exists only during daemon authentication. Polling the
78    // authoritative connection map avoids retaining a detached waiter after a
79    // socket disconnect without adding a second per-connection ownership map.
80    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        // C++ `APIServer_InitConnect.cpp` special-cases exactly this internal
158        // UI client and replies immediately; all ordinary clients wait for the
159        // gateway init result instead of receiving a successful uid=0.
160        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    /// Request closure exactly once.
251    ///
252    /// The return value describes only the `open -> closing` transition.
253    /// Tokio permits updating a watch value after all receivers are dropped;
254    /// receiver presence therefore intentionally does not affect the result.
255    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    /// 生成随机连接 ID(与 C++ 的 GetRand_MilliTimeAndU22 对应)
420    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    /// 生成随机 AES key(16 字节 hex 字符串的 ASCII 字节)
435    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    /// 创建发送帧,自动处理 AES 加密
444    ///
445    /// 当 aes_encrypt_enabled 为 true 时:
446    /// - SHA1 基于明文计算
447    /// - body 使用 AES-128 ECB 加密
448    /// - header.body_len 更新为密文长度
449    ///
450    /// 对应 C++ APIServerCS_Conn::OnSendPacketData 的加密逻辑
451    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    /// 创建发送帧,复用调用方已计算的明文 body SHA1。
457    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    /// 解密请求 body(如果启用了 AES 加密)
473    ///
474    /// 对应 C++ APIServerCS_Conn::OnRecvPacket 的解密逻辑
475    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    /// 处理 InitConnect 请求,返回 InitConnect 响应 body
487    ///
488    /// 当配置了 RSA 私钥时:
489    /// - C2S 请求 body 使用 RSA 公钥加密(需要用私钥解密)
490    /// - S2C 响应 body 使用 RSA 公钥加密(客户端用私钥解密)
491    ///
492    /// 对应 C++ APIServer::OnRecvInitConnect
493    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        // 1. 解密 C2S(如果配置了 RSA)
536        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        // 2. 加密 S2C(如果配置了 RSA)
648        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    /// 处理 KeepAlive 请求。
666    ///
667    /// This compatibility helper falls back to the local daemon clock. The
668    /// TCP / WebSocket dispatch paths must inject the backend-adjusted server
669    /// time through [`Self::handle_keepalive_at`] to match C++ OpenD.
670    pub fn handle_keepalive(&self, body: &[u8]) -> Result<Vec<u8>, FutuError> {
671        self.handle_keepalive_at(body, chrono::Utc::now().timestamp())
672    }
673
674    /// 处理 KeepAlive 请求,使用调用方注入的 server time。
675    ///
676    /// Ref: C++ `APIServerCS_Conn.cpp:370-373` replies with
677    /// `INNBiz_SvrTime::GetSvrTimeStamp()`.
678    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
701/// 连接断开通知
702pub struct DisconnectNotify {
703    /// 被断开的连接 ID(订阅 / push / auth 状态清理用)
704    pub conn_id: u64,
705}
706
707/// 通知 listener 清理连接;cleanup task 已退出时至少留下可观测日志。
708pub(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
723/// 运行单个连接的收发循环。
724///
725/// 保留 v1.4.x 的 public 返回形状;server listener 使用 crate-private sibling
726/// 取得额外的 close control。
727pub 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    // 发送任务
784    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    // 接收任务
797    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;