1use std::collections::HashMap;
8use std::sync::Arc;
9use std::sync::atomic::{AtomicU32, Ordering};
10use std::time::SystemTime;
11
12use arc_swap::ArcSwapOption;
13use bytes::Bytes;
14use futures::{SinkExt, StreamExt};
15use parking_lot::Mutex;
16use tokio::net::TcpStream;
17use tokio::sync::{mpsc, oneshot, watch};
18use tokio_util::codec::Framed;
19
20use futu_core::error::{FutuError, Result};
21use futu_core::log_redact::endpoint_log_fingerprint;
22use futu_net::encrypt::{aes_cbc_md5_decrypt_var, aes_cbc_md5_encrypt_var};
23
24use crate::nn_codec::{ExHeadErrorInfo, NNCodec, NNFrame, NNHeader, should_skip_encryption};
25
26pub struct BackendConn {
28 serial_no: AtomicU32,
29 sec_data: AtomicU32,
30 connected: Arc<std::sync::atomic::AtomicBool>,
31 connected_tx: watch::Sender<bool>,
32 session_key: SharedSessionKey,
38 cmd_tx: mpsc::Sender<BackendCmd>,
39 pending: PendingResponses,
40 shutdown_tx: watch::Sender<bool>,
41 pub user_id: AtomicU32,
42 client_ip: Mutex<String>,
47 pub client_type: u8,
48 pub client_ver: u16,
49 pub lang_id: u8,
50}
51
52enum BackendCmd {
53 Send(NNFrame),
54}
55
56type PendingResponses = Arc<Mutex<HashMap<u32, oneshot::Sender<NNFrame>>>>;
57type SharedSessionKey = Arc<ArcSwapOption<Vec<u8>>>;
58
59pub type PushCallback = Arc<dyn Fn(u16, Bytes, SystemTime) + Send + Sync + 'static>;
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63enum InboundFrameDecision {
64 Deliver,
65 Drop,
66}
67
68fn decode_inbound_frame_body(
69 frame: &mut NNFrame,
70 session_key: Option<&[u8]>,
71) -> InboundFrameDecision {
72 if !should_skip_encryption(frame.header.cmd_id)
73 && let Some(key) = session_key
74 {
75 let body_len = frame.body.len();
76 if body_len >= 32 && body_len.is_multiple_of(16) {
80 match aes_cbc_md5_decrypt_var(key, &frame.body) {
81 Ok(decrypted) => {
82 frame.body = Bytes::from(decrypted);
84 }
85 Err(e) => {
86 tracing::warn!(
90 cmd_id = frame.header.cmd_id,
91 body_len = body_len,
92 key_len = key.len(),
93 error = %e,
94 "decrypt failed, dropping inbound frame"
95 );
96 return InboundFrameDecision::Drop;
97 }
98 }
99 } else {
100 tracing::debug!(
101 cmd_id = frame.header.cmd_id,
102 body_len = body_len,
103 "body not encrypted (len not aligned to 16)"
104 );
105 }
106 }
107
108 if frame.header.is_compressed() {
109 let compressed_body_len = frame.body.len();
110 match crate::ftlogin_wire::decode_inbound_body(true, frame.body.as_ref()) {
111 Ok(decompressed) => {
112 frame.body = Bytes::from(decompressed);
113 frame.header.body_len = frame.body.len() as u32;
114 tracing::debug!(
115 cmd_id = frame.header.cmd_id,
116 serial_no = frame.header.serial_no,
117 compressed_body_len,
118 body_len = frame.body.len(),
119 "decompressed inbound frame after decrypt"
120 );
121 }
122 Err(e) => {
123 tracing::warn!(
128 cmd_id = frame.header.cmd_id,
129 serial_no = frame.header.serial_no,
130 body_len = compressed_body_len,
131 error = %e,
132 "decompress failed after decrypt, dropping inbound frame"
133 );
134 return InboundFrameDecision::Drop;
135 }
136 }
137 }
138
139 InboundFrameDecision::Deliver
140}
141
142fn release_pending_on_inbound_decode_error(frame: &NNFrame, pending: &PendingResponses) {
143 if !frame.header.is_push {
144 pending.lock().remove(&frame.header.serial_no);
145 }
146}
147
148fn is_ex_head_error(err: &ExHeadErrorInfo) -> bool {
149 err.cmd_result != 0 || err.code != 0 || !err.message.is_empty() || !err.source.is_empty()
150}
151
152fn response_ex_head_error(frame: &NNFrame) -> Option<FutuError> {
153 let err = frame.parse_ex_head_error()?;
154 if !is_ex_head_error(&err) {
155 return None;
156 }
157
158 let ret_type = if err.code != 0 {
159 err.code
160 } else {
161 err.cmd_result
162 };
163 Some(FutuError::ServerError {
164 ret_type,
165 msg: format!(
166 "backend cmd {} ex_head error: cmd_result={} code={} source={} message={}",
167 frame.header.cmd_id, err.cmd_result, err.code, err.source, err.message
168 ),
169 })
170}
171
172fn mark_disconnected(
173 connected: &Arc<std::sync::atomic::AtomicBool>,
174 connected_tx: &watch::Sender<bool>,
175) {
176 let was_connected = connected.swap(false, Ordering::AcqRel);
177 if was_connected {
178 let _ = connected_tx.send(false);
179 }
180}
181
182impl BackendConn {
183 pub const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
190
191 pub const CLIENT_VER_FTGTW: u16 = 1030;
198
199 async fn establish_stream(addr: &str, timeout: std::time::Duration) -> Result<TcpStream> {
201 let stream = match tokio::time::timeout(timeout, TcpStream::connect(addr)).await {
202 Ok(Ok(s)) => s,
203 Ok(Err(e)) => return Err(e.into()),
204 Err(_elapsed) => {
205 return Err(FutuError::Network(std::io::Error::new(
206 std::io::ErrorKind::TimedOut,
207 format!("connect to {addr} timed out after {}s", timeout.as_secs()),
208 )));
209 }
210 };
211 stream.set_nodelay(true)?;
212 Ok(stream)
213 }
214
215 pub async fn connect(addr: &str, push_callback: PushCallback) -> Result<Self> {
217 let stream = Self::establish_stream(addr, Self::CONNECT_TIMEOUT).await?;
218 tracing::info!(
219 addr_fp = %endpoint_log_fingerprint(addr),
220 "connected to backend"
221 );
222 Ok(Self::from_stream(stream, push_callback))
223 }
224
225 pub async fn connect_race(
235 addrs: &[String],
236 push_callback: PushCallback,
237 ) -> Result<(Self, String)> {
238 use futures::stream::{FuturesUnordered, StreamExt};
239
240 if addrs.is_empty() {
241 return Err(FutuError::Network(std::io::Error::new(
242 std::io::ErrorKind::InvalidInput,
243 "connect_race: empty address list",
244 )));
245 }
246
247 tracing::info!(
248 candidates = addrs.len(),
249 candidate_fps = ?addrs
250 .iter()
251 .map(|addr| endpoint_log_fingerprint(addr))
252 .collect::<Vec<_>>(),
253 "racing parallel connects"
254 );
255
256 let mut attempts: FuturesUnordered<_> = addrs
257 .iter()
258 .cloned()
259 .map(|addr| async move {
260 let result = Self::establish_stream(&addr, Self::CONNECT_TIMEOUT).await;
261 (addr, result)
262 })
263 .collect();
264
265 let mut last_err: Option<FutuError> = None;
266 while let Some((addr, result)) = attempts.next().await {
267 match result {
268 Ok(stream) => {
269 tracing::info!(
270 addr_fp = %endpoint_log_fingerprint(&addr),
271 remaining_losers = attempts.len(),
272 "connect race winner"
273 );
274 drop(attempts); let conn = Self::from_stream(stream, push_callback);
276 return Ok((conn, addr));
277 }
278 Err(e) => {
279 tracing::debug!(
280 addr_fp = %endpoint_log_fingerprint(&addr),
281 error = %e,
282 "candidate failed"
283 );
284 last_err = Some(e);
285 }
286 }
287 }
288
289 Err(last_err.unwrap_or_else(|| {
290 FutuError::Network(std::io::Error::other("connect_race: all candidates failed"))
291 }))
292 }
293
294 #[cfg(feature = "test-util")]
299 pub fn from_duplex(stream: tokio::io::DuplexStream, push_callback: PushCallback) -> Self {
300 Self::from_stream_inner(stream, push_callback)
301 }
302
303 fn from_stream(stream: TcpStream, push_callback: PushCallback) -> Self {
305 Self::from_stream_inner(stream, push_callback)
306 }
307
308 pub(crate) fn from_stream_inner<S>(stream: S, push_callback: PushCallback) -> Self
311 where
312 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static,
313 {
314 let framed = Framed::new(stream, NNCodec);
315 let (mut sink, mut stream_rx) = framed.split();
316
317 let (cmd_tx, mut cmd_rx) = mpsc::channel::<BackendCmd>(256);
318
319 let pending: PendingResponses =
320 Arc::new(Mutex::new(HashMap::<u32, oneshot::Sender<NNFrame>>::new()));
321 let pending_recv = pending.clone();
322 let pending_send = pending.clone();
323 let connected = Arc::new(std::sync::atomic::AtomicBool::new(true));
324 let (connected_tx, _) = watch::channel(true);
325 let (shutdown_tx, mut shutdown_rx_recv) = watch::channel(false);
326 let mut shutdown_rx_send = shutdown_tx.subscribe();
327 let connected_recv = Arc::clone(&connected);
328 let connected_send = Arc::clone(&connected);
329 let connected_tx_recv = connected_tx.clone();
330 let connected_tx_send = connected_tx.clone();
331 let session_key: SharedSessionKey = Arc::new(ArcSwapOption::empty());
332 let session_key_for_recv = session_key.clone();
333
334 tokio::spawn(async move {
336 loop {
337 let next_frame = tokio::select! {
338 _ = shutdown_rx_recv.changed() => {
339 break;
340 }
341 next = stream_rx.next() => next,
342 };
343 let Some(Ok(mut frame)) = next_frame else {
344 break;
345 };
346
347 tracing::debug!(
349 cmd_id = frame.header.cmd_id,
350 serial_no = frame.header.serial_no,
351 is_push = frame.header.is_push,
352 is_compressed = frame.header.is_compressed,
353 ex_head_len = frame.header.ex_head_len,
354 body_len = frame.header.body_len,
355 actual_body_len = frame.body.len(),
356 "recv frame"
357 );
358
359 if let Some(err) = frame.parse_ex_head_error()
367 && is_ex_head_error(&err)
368 {
369 let is_soft_fail = err.code == -102;
375 if is_soft_fail {
376 tracing::debug!(
377 cmd_id = frame.header.cmd_id,
378 code = err.code,
379 source = %err.source,
380 message = %err.message,
381 "server: cmd not available on this channel (soft-fail)"
382 );
383 } else {
384 tracing::warn!(
385 cmd_id = frame.header.cmd_id,
386 cmd_result = err.cmd_result,
387 code = err.code,
388 source = %err.source,
389 message = %err.message,
390 "server returned err_info in ex_head"
391 );
392 }
393 }
394
395 let session_key = session_key_for_recv.load_full();
396 if decode_inbound_frame_body(&mut frame, session_key.as_deref().map(Vec::as_slice))
397 == InboundFrameDecision::Drop
398 {
399 release_pending_on_inbound_decode_error(&frame, &pending_recv);
400 continue;
401 }
402
403 let is_push = frame.header.is_push
408 || (frame.header.serial_no == 0 && !pending_recv.lock().contains_key(&0));
409 if is_push {
410 tracing::debug!(
411 cmd_id = frame.header.cmd_id,
412 body_len = frame.body.len(),
413 is_push = frame.header.is_push,
414 is_compressed = frame.header.is_compressed,
415 reserved = ?frame.header.reserved,
416 "backend push received"
417 );
418 push_callback(frame.header.cmd_id, frame.body, SystemTime::now());
419 } else {
420 let tx = pending_recv.lock().remove(&frame.header.serial_no);
421 if let Some(tx) = tx
422 && let Err(frame) = tx.send(frame)
423 {
424 tracing::debug!(
425 cmd_id = frame.header.cmd_id,
426 serial_no = frame.header.serial_no,
427 body_len = frame.body.len(),
428 "backend response receiver dropped before frame delivery"
429 );
430 }
431 }
432 }
433 mark_disconnected(&connected_recv, &connected_tx_recv);
435 tracing::warn!("backend connection closed");
436 let mut pending = pending_recv.lock();
437 let count = pending.len();
438 if count > 0 {
439 tracing::warn!(
440 pending_count = count,
441 "aborting pending requests due to disconnect"
442 );
443 }
444 pending.clear();
448 });
449
450 tokio::spawn(async move {
452 loop {
453 let next_cmd = tokio::select! {
454 _ = shutdown_rx_send.changed() => {
455 break;
456 }
457 cmd = cmd_rx.recv() => cmd,
458 };
459 let Some(cmd) = next_cmd else {
460 break;
461 };
462
463 match cmd {
464 BackendCmd::Send(frame) => {
465 if let Err(e) = sink.send(frame).await {
466 mark_disconnected(&connected_send, &connected_tx_send);
467 tracing::error!(error = %e, "backend send failed");
468 let mut pending = pending_send.lock();
469 let count = pending.len();
470 if count > 0 {
471 tracing::warn!(
472 pending_count = count,
473 "aborting pending requests due to send failure"
474 );
475 }
476 pending.clear();
477 break;
478 }
479 }
480 }
481 }
482 });
483
484 Self {
485 serial_no: AtomicU32::new(0),
486 sec_data: AtomicU32::new(1),
487 connected,
488 connected_tx,
489 session_key, cmd_tx,
491 pending,
492 shutdown_tx,
493 user_id: AtomicU32::new(0),
494 client_ip: Mutex::new(String::new()),
495 client_type: 40, client_ver: Self::CLIENT_VER_FTGTW,
497 lang_id: 0,
498 }
499 }
500
501 pub fn set_session_key(&self, key: Vec<u8>) {
505 self.session_key.store(Some(Arc::new(key)));
506 }
507
508 pub fn set_sec_data(&self, val: u32) {
510 self.sec_data.store(val, Ordering::Relaxed);
511 }
512
513 pub fn set_client_ip(&self, ip: String) {
515 *self.client_ip.lock() = ip;
516 }
517
518 pub fn client_ip(&self) -> String {
520 self.client_ip.lock().clone()
521 }
522
523 pub fn is_connected(&self) -> bool {
524 self.connected.load(Ordering::Acquire)
525 }
526
527 pub fn subscribe_connection_state(&self) -> watch::Receiver<bool> {
535 self.connected_tx.subscribe()
536 }
537
538 pub fn mark_disconnected_for_reconnect(&self, reason: &'static str) {
546 tracing::warn!(
547 reason,
548 "marking backend connection disconnected to trigger reconnect"
549 );
550 mark_disconnected(&self.connected, &self.connected_tx);
551 }
552
553 pub async fn request(&self, cmd_id: u16, body: Vec<u8>) -> Result<NNFrame> {
555 self.request_with_reserved(cmd_id, body, [0u8; 10]).await
556 }
557
558 pub async fn request_with_reserved(
560 &self,
561 cmd_id: u16,
562 body: Vec<u8>,
563 reserved: [u8; 10],
564 ) -> Result<NNFrame> {
565 self.request_with_reserved_timeout(
566 cmd_id,
567 body,
568 reserved,
569 std::time::Duration::from_secs(10),
570 )
571 .await
572 }
573
574 pub(crate) async fn request_with_reserved_timeout(
575 &self,
576 cmd_id: u16,
577 body: Vec<u8>,
578 reserved: [u8; 10],
579 timeout: std::time::Duration,
580 ) -> Result<NNFrame> {
581 if !self.is_connected() {
582 return Err(FutuError::NotInitialized);
583 }
584
585 let deadline = tokio::time::Instant::now() + timeout;
586 let frame = self.build_outbound_frame(cmd_id, body, reserved)?;
587 let serial_no = frame.header.serial_no;
588
589 let (resp_tx, resp_rx) = oneshot::channel();
590 self.pending.lock().insert(serial_no, resp_tx);
591 match tokio::time::timeout_at(deadline, self.cmd_tx.send(BackendCmd::Send(frame))).await {
592 Ok(Ok(())) => {}
593 Ok(Err(_closed)) => {
594 self.pending.lock().remove(&serial_no);
595 mark_disconnected(&self.connected, &self.connected_tx);
596 return Err(FutuError::NotInitialized);
597 }
598 Err(_elapsed) => {
599 self.pending.lock().remove(&serial_no);
600 mark_disconnected(&self.connected, &self.connected_tx);
601 return Err(FutuError::Timeout);
602 }
603 }
604
605 let resp = crate::delay_stats::trace_backend_request(cmd_id, async {
606 match tokio::time::timeout_at(deadline, resp_rx).await {
607 Ok(Ok(resp)) => Ok(resp),
608 Ok(Err(_closed)) => {
609 self.pending.lock().remove(&serial_no);
610 Err(FutuError::Codec("response channel closed".into()))
611 }
612 Err(_elapsed) => {
613 self.pending.lock().remove(&serial_no);
614 mark_disconnected(&self.connected, &self.connected_tx);
615 Err(FutuError::Timeout)
616 }
617 }
618 })
619 .await?;
620
621 if let Some(err) = response_ex_head_error(&resp) {
622 return Err(err);
623 }
624
625 Ok(resp)
626 }
627
628 pub async fn send_fire_and_forget(&self, cmd_id: u16, body: Vec<u8>) -> Result<()> {
630 let frame = self.build_outbound_frame(cmd_id, body, [0u8; 10])?;
631 if self.cmd_tx.send(BackendCmd::Send(frame)).await.is_err() {
632 mark_disconnected(&self.connected, &self.connected_tx);
633 return Err(FutuError::NotInitialized);
634 }
635
636 Ok(())
637 }
638
639 fn build_outbound_frame(
640 &self,
641 cmd_id: u16,
642 body: Vec<u8>,
643 reserved: [u8; 10],
644 ) -> Result<NNFrame> {
645 let serial = self.next_serial();
646 let mut header = NNHeader::new(cmd_id, serial);
647 header.user_id = self.user_id.load(Ordering::Relaxed);
648 header.client_type = self.client_type;
649 header.client_ver = self.client_ver;
650 header.lang_id = self.lang_id;
651 header.reserved.copy_from_slice(&reserved[..8]);
655
656 let final_body = self.encode_outbound_body(cmd_id, body)?;
657 header.body_len = final_body.len() as u32;
658
659 Ok(NNFrame {
660 header,
661 body: Bytes::from(final_body),
662 ex_head: Bytes::new(),
663 })
664 }
665
666 fn encode_outbound_body(&self, cmd_id: u16, body: Vec<u8>) -> Result<Vec<u8>> {
667 if should_skip_encryption(cmd_id) {
668 return Ok(body);
669 }
670
671 let key = self.session_key.load_full();
672 match key {
673 Some(key) => {
674 let sec = self.sec_data.fetch_add(1, Ordering::Relaxed) + 1;
676 let mut plaintext = Vec::with_capacity(4 + body.len());
677 plaintext.extend_from_slice(&sec.to_be_bytes());
678 plaintext.extend_from_slice(&body);
679 aes_cbc_md5_encrypt_var(key.as_slice(), &plaintext)
681 }
682 None => Ok(body),
683 }
684 }
685
686 fn next_serial(&self) -> u32 {
687 self.serial_no.fetch_add(1, Ordering::Relaxed) + 1
688 }
689}
690
691impl Drop for BackendConn {
692 fn drop(&mut self) {
693 let _ = self.shutdown_tx.send(true);
694 mark_disconnected(&self.connected, &self.connected_tx);
695 self.pending.lock().clear();
696 }
697}
698
699#[cfg(test)]
700mod tests;