futu_net/

client.rs

use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Duration;

use bytes::Bytes;
use dashmap::DashMap;
use tokio::sync::{mpsc, oneshot};

use futu_codec::frame::FutuFrame;
use futu_core::error::{FutuError, Result};
use futu_core::proto_id;

use crate::connection::Connection;
use crate::encrypt;
use crate::reconnect::ReconnectPolicy;

const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(12);
const DEFAULT_KEEP_ALIVE_INTERVAL_SECS: u64 = 10;
const MIN_KEEP_ALIVE_INTERVAL_SECS: u64 = 1;
const MAX_KEEP_ALIVE_INTERVAL_SECS: u64 = 60;
const MAX_HEARTBEAT_SEND_FAILURES: u32 = 3;
pub(crate) const PUSH_CHANNEL_CAPACITY: usize = 4096;

fn sanitize_keep_alive_interval_secs(raw: i32) -> u64 {
    if raw <= 0 {
        return DEFAULT_KEEP_ALIVE_INTERVAL_SECS;
    }
    (raw as u64).clamp(MIN_KEEP_ALIVE_INTERVAL_SECS, MAX_KEEP_ALIVE_INTERVAL_SECS)
}

fn heartbeat_send_failure_should_exit(consecutive_failures: &mut u32, result: &Result<()>) -> bool {
    if result.is_ok() {
        *consecutive_failures = 0;
        return false;
    }

    *consecutive_failures = consecutive_failures.saturating_add(1);
    *consecutive_failures >= MAX_HEARTBEAT_SEND_FAILURES
}

fn init_connect_server_err(ret_type: i32, ret_msg: Option<String>) -> FutuError {
    let msg = match ret_msg {
        Some(msg) if !msg.is_empty() => msg,
        _ => "<missing ret_msg>".to_string(),
    };
    FutuError::ServerError { ret_type, msg }
}

/// 客户端配置
#[derive(Debug, Clone)]
pub struct ClientConfig {
    pub addr: String,
    pub client_ver: String,
    pub client_id: String,
    pub recv_notify: bool,
    /// RSA 私钥（PEM 格式）。若提供则启用加密：
    /// - InitConnect 使用 RSA 加解密
    /// - 后续请求使用 AES-128 ECB 加解密
    /// - 若不提供则全程明文通信
    pub rsa_key: Option<String>,
}

/// 推送消息
#[derive(Debug, Clone)]
pub struct PushMessage {
    pub proto_id: u32,
    pub body: Bytes,
}

type PushSender = mpsc::Sender<PushMessage>;
pub type PushReceiver = mpsc::Receiver<PushMessage>;

/// 请求上下文（等待响应的 oneshot sender）
struct PendingRequest {
    tx: oneshot::Sender<FutuFrame>,
}

/// FutuOpenD 高层客户端
///
/// 功能：
/// - InitConnect 握手
/// - 自动心跳 KeepAlive
/// - 请求-响应匹配（serial number）
/// - 推送消息分发
/// - 断线重连
pub struct FutuClient {
    config: ClientConfig,
    serial_no: AtomicU32,
    pending: Arc<DashMap<u32, PendingRequest>>,
    push_tx: PushSender,
    cmd_tx: Option<mpsc::Sender<ClientCommand>>,
    conn_aes_key: parking_lot::Mutex<Option<[u8; 16]>>,
    conn_id: AtomicU64,
    init_ai_type: Option<i32>,
}

enum ClientCommand {
    Send(FutuFrame, oneshot::Sender<()>),
}

fn decode_client_inbound_body(frame: &FutuFrame, aes_key: Option<&[u8; 16]>) -> Result<Bytes> {
    let body = match aes_key {
        Some(key) if !frame.body.is_empty() => {
            encrypt::aes_ecb_decrypt(key, &frame.body).map(Bytes::from)
        }
        _ => Ok(frame.body.clone()),
    }?;
    verify_plaintext_body_sha1(frame, body)
}

fn verify_plaintext_body_sha1(frame: &FutuFrame, plaintext: Bytes) -> Result<Bytes> {
    let plaintext_frame = FutuFrame {
        header: frame.header.clone(),
        body: plaintext.clone(),
    };
    if plaintext_frame.verify_sha1() {
        Ok(plaintext)
    } else {
        Err(FutuError::Sha1Mismatch)
    }
}

fn build_init_connect_request(
    config: &ClientConfig,
    ai_type: Option<i32>,
) -> futu_proto::init_connect::Request {
    futu_proto::init_connect::Request {
        c2s: futu_proto::init_connect::C2s {
            client_ver: 100,
            client_id: config.client_id.clone(),
            recv_notify: Some(config.recv_notify),
            packet_enc_algo: Some(0),
            push_proto_fmt: Some(0),
            programming_language: Some(String::new()),
            ai_type,
        },
    }
}

fn parse_conn_aes_key(conn_aes_key: &str) -> Result<[u8; 16]> {
    let key_bytes = conn_aes_key.as_bytes();
    match key_bytes.len() {
        16 => {
            let mut key = [0u8; 16];
            key.copy_from_slice(key_bytes);
            Ok(key)
        }
        32 => hex_decode_16(key_bytes)
            .ok_or_else(|| FutuError::Codec("invalid hex InitConnect AES key".into())),
        len => Err(FutuError::Codec(format!(
            "unexpected AES key length: {len}"
        ))),
    }
}

fn release_pending_on_inbound_decode_error(
    frame: &FutuFrame,
    pending: &DashMap<u32, PendingRequest>,
) {
    if !proto_id::is_push_proto(frame.header.proto_id) {
        pending.remove(&frame.header.serial_no);
    }
}

fn push_channel() -> (PushSender, PushReceiver) {
    mpsc::channel(PUSH_CHANNEL_CAPACITY)
}

fn send_push_message(push_tx: &PushSender, message: PushMessage) -> bool {
    let proto_id = message.proto_id;
    match push_tx.try_send(message) {
        Ok(()) => true,
        Err(mpsc::error::TrySendError::Full(_)) => {
            tracing::warn!(
                proto_id,
                capacity = PUSH_CHANNEL_CAPACITY,
                "push receiver is slow; dropping push message from bounded client queue"
            );
            false
        }
        Err(mpsc::error::TrySendError::Closed(_)) => {
            tracing::debug!(proto_id, "push receiver dropped before delivery");
            false
        }
    }
}

fn send_response_frame(tx: oneshot::Sender<FutuFrame>, frame: FutuFrame) -> bool {
    let serial = frame.header.serial_no;
    let proto_id = frame.header.proto_id;
    match tx.send(frame) {
        Ok(()) => true,
        Err(_) => {
            tracing::debug!(
                serial,
                proto_id,
                "response receiver dropped before delivery"
            );
            false
        }
    }
}

fn send_command_ack(ack: oneshot::Sender<()>) -> bool {
    match ack.send(()) {
        Ok(()) => true,
        Err(_) => {
            tracing::debug!("command ack receiver dropped before delivery");
            false
        }
    }
}

impl FutuClient {
    /// 创建客户端（未连接状态）
    ///
    /// 返回 (client, push_rx)，其中 push_rx 用于接收服务端推送。
    pub fn new(config: ClientConfig) -> (Self, PushReceiver) {
        let (push_tx, push_rx) = push_channel();

        let client = Self {
            config,
            serial_no: AtomicU32::new(0),
            pending: Arc::new(DashMap::new()),
            push_tx,
            cmd_tx: None,
            conn_aes_key: parking_lot::Mutex::new(None),
            conn_id: AtomicU64::new(0),
            init_ai_type: None,
        };

        (client, push_rx)
    }

    /// Set C++ 10.7 `InitConnect.C2S.aiType` for the next connection.
    ///
    /// `None` preserves the legacy wire shape. C++ treats an absent value as
    /// `0` (non-AI); `Some(1)` marks a skills/AI caller.
    pub fn with_init_ai_type(mut self, ai_type: Option<i32>) -> Self {
        self.init_ai_type = ai_type;
        self
    }

    /// 连接到 OpenD 网关并完成 InitConnect 握手
    pub async fn connect(&mut self) -> Result<InitConnectInfo> {
        let mut conn = Connection::connect(&self.config.addr).await?;

        // 发送 InitConnect 请求
        let serial = self.next_serial();
        let req = build_init_connect_request(&self.config, self.init_ai_type);
        let raw_body = prost::Message::encode_to_vec(&req);

        // 如果有 RSA 密钥，用 RSA 公钥加密 InitConnect 请求 body
        let send_body = if let Some(ref rsa_key) = self.config.rsa_key {
            tracing::debug!("encrypting InitConnect with RSA");
            encrypt::rsa_public_encrypt(rsa_key, &raw_body)?
        } else {
            raw_body.clone()
        };

        // SHA1 基于明文
        let mut frame = Connection::build_frame(proto_id::INIT_CONNECT, serial, send_body);
        {
            use sha1::{Digest, Sha1};
            let mut hasher = Sha1::new();
            hasher.update(&raw_body);
            let hash = hasher.finalize();
            frame.header.body_sha1.copy_from_slice(&hash);
        }
        conn.send(frame).await?;

        // 接收 InitConnect 响应
        let resp_frame = conn.recv().await?.ok_or(FutuError::Codec(
            "connection closed during InitConnect".into(),
        ))?;

        // 如果有 RSA 密钥，用 RSA 私钥解密 InitConnect 响应 body
        let resp_body = if let Some(ref rsa_key) = self.config.rsa_key {
            tracing::debug!("decrypting InitConnect response with RSA");
            Bytes::from(encrypt::rsa_private_decrypt(rsa_key, &resp_frame.body)?)
        } else {
            resp_frame.body.clone()
        };
        let resp_body = verify_plaintext_body_sha1(&resp_frame, resp_body)?;

        let resp: futu_proto::init_connect::Response =
            prost::Message::decode(resp_body.as_ref()).map_err(FutuError::Proto)?;

        // 检查返回码
        let ret_type = resp.ret_type;
        if ret_type != 0 {
            return Err(init_connect_server_err(ret_type, resp.ret_msg));
        }

        let s2c = resp.s2c.ok_or(FutuError::Codec(
            "missing s2c in InitConnect response".into(),
        ))?;

        let info = InitConnectInfo {
            server_ver: s2c.server_ver,
            login_user_id: s2c.login_user_id,
            conn_id: s2c.conn_id,
            conn_aes_key: s2c.conn_aes_key.clone(),
            keep_alive_interval: s2c.keep_alive_interval,
        };
        self.conn_id.store(info.conn_id, Ordering::Relaxed);

        // 保存 AES key
        if !info.conn_aes_key.is_empty() {
            let key_bytes = info.conn_aes_key.as_bytes();
            tracing::debug!(key_len = key_bytes.len(), "received AES key");
            *self.conn_aes_key.lock() = Some(parse_conn_aes_key(&info.conn_aes_key)?);
        }

        tracing::info!(
            server_ver = info.server_ver,
            conn_id = info.conn_id,
            keep_alive_interval = info.keep_alive_interval,
            "InitConnect succeeded"
        );

        // 启动后台任务：心跳、消息接收。C++ OpenD API server 固定下发 10s，
        // 且每 66s 检查活跃计数；这里对异常 S2C 做本地边界保护，避免
        // malicious/misconfigured server 传 0 或极大值导致 tokio interval panic
        // 或长时间不发心跳。
        let keep_alive_interval_secs = sanitize_keep_alive_interval_secs(info.keep_alive_interval);
        if info.keep_alive_interval <= 0
            || keep_alive_interval_secs != info.keep_alive_interval as u64
        {
            tracing::warn!(
                raw_keep_alive_interval = info.keep_alive_interval,
                sanitized_keep_alive_interval = keep_alive_interval_secs,
                "sanitized InitConnect keep_alive_interval"
            );
        }
        let keep_alive_interval = Duration::from_secs(keep_alive_interval_secs);
        self.start_background_tasks(conn, keep_alive_interval);

        Ok(info)
    }

    /// 发送请求并等待响应
    pub async fn request(&self, proto_id: u32, body: Vec<u8>) -> Result<FutuFrame> {
        self.request_with_timeout(proto_id, body, DEFAULT_REQUEST_TIMEOUT)
            .await
    }

    /// 发送请求并等待响应（调用方可指定总请求等待上限）。
    ///
    /// 与 [`Self::request`] 共用同一 pending 清理路径：command channel 入队、
    /// send ack、响应等待任一阶段超时时都必须移除 `pending[serial]`，避免
    /// 调用方自己在外层 `tokio::time::timeout` 取消 future 时留下悬挂请求。
    pub async fn request_with_timeout(
        &self,
        proto_id: u32,
        body: Vec<u8>,
        timeout: Duration,
    ) -> Result<FutuFrame> {
        let deadline = tokio::time::Instant::now() + timeout;
        let serial = self.next_serial();

        // 加密 body（如果有 AES key）
        let (final_body, sha1) = self.prepare_body(&body);

        let frame = FutuFrame::with_sha1(proto_id, serial, Bytes::from(final_body), sha1);

        let (resp_tx, resp_rx) = oneshot::channel();
        self.pending.insert(serial, PendingRequest { tx: resp_tx });

        // 通过 command channel 发送
        if let Some(cmd_tx) = &self.cmd_tx {
            let (ack_tx, ack_rx) = oneshot::channel();
            match tokio::time::timeout_at(deadline, cmd_tx.send(ClientCommand::Send(frame, ack_tx)))
                .await
            {
                Ok(Ok(())) => {}
                Ok(Err(_)) => {
                    self.pending.remove(&serial);
                    return Err(FutuError::NotInitialized);
                }
                Err(_) => {
                    self.pending.remove(&serial);
                    return Err(FutuError::Timeout);
                }
            }

            match tokio::time::timeout_at(deadline, ack_rx).await {
                Ok(Ok(())) => {}
                Ok(Err(_)) => {
                    self.pending.remove(&serial);
                    return Err(FutuError::Codec("send ack failed".into()));
                }
                Err(_) => {
                    self.pending.remove(&serial);
                    return Err(FutuError::Timeout);
                }
            }
        } else {
            self.pending.remove(&serial);
            return Err(FutuError::NotInitialized);
        }

        // 等待响应（带超时）
        let resp = tokio::time::timeout_at(deadline, resp_rx)
            .await
            .map_err(|elapsed| {
                self.pending.remove(&serial);
                tracing::warn!(
                    serial,
                    error = %elapsed,
                    "request response wait timed out; pending entry removed"
                );
                FutuError::Timeout
            })?
            .map_err(|err| {
                FutuError::Codec(format!(
                    "response channel closed for serial {serial}: {err}"
                ))
            })?;

        Ok(resp)
    }

    /// Server-assigned connection id from InitConnect S2C.
    ///
    /// FTAPI trade-write packet ids must echo this value in `PacketID.connID`;
    /// the gateway replay guard compares it against the actual TCP connection.
    pub fn conn_id(&self) -> Option<u64> {
        let conn_id = self.conn_id.load(Ordering::Relaxed);
        (conn_id != 0).then_some(conn_id)
    }

    fn next_serial(&self) -> u32 {
        self.serial_no.fetch_add(1, Ordering::Relaxed) + 1
    }

    fn prepare_body(&self, plaintext: &[u8]) -> (Vec<u8>, [u8; 20]) {
        use sha1::{Digest, Sha1};

        // SHA1 始终基于明文
        let mut hasher = Sha1::new();
        hasher.update(plaintext);
        let sha1_result = hasher.finalize();
        let mut sha1 = [0u8; 20];
        sha1.copy_from_slice(&sha1_result);

        // 仅在有 RSA 密钥（即启用加密）时才用 AES 加密
        let body = if self.config.rsa_key.is_some() {
            let key = self.conn_aes_key.lock();
            match key.as_ref() {
                Some(k) => encrypt::aes_ecb_encrypt(k, plaintext),
                None => plaintext.to_vec(),
            }
        } else {
            plaintext.to_vec()
        };

        (body, sha1)
    }

    fn start_background_tasks(&mut self, conn: Connection, keep_alive_interval: Duration) {
        let (cmd_tx, cmd_rx) = mpsc::channel(256);
        self.cmd_tx = Some(cmd_tx.clone());

        let pending = Arc::clone(&self.pending);
        let push_tx = self.push_tx.clone();
        let aes_key = if self.config.rsa_key.is_some() {
            *self.conn_aes_key.lock()
        } else {
            None
        };

        tokio::spawn(async move {
            run_event_loop(conn, cmd_rx, pending, push_tx, aes_key, keep_alive_interval).await;
        });
    }
}

/// 后台事件循环：处理接收、心跳、发送
async fn run_event_loop(
    mut conn: Connection,
    mut cmd_rx: mpsc::Receiver<ClientCommand>,
    pending: Arc<DashMap<u32, PendingRequest>>,
    push_tx: PushSender,
    aes_key: Option<[u8; 16]>,
    keep_alive_interval: Duration,
) {
    let mut heartbeat = tokio::time::interval(keep_alive_interval);
    heartbeat.tick().await; // 跳过第一次立即触发
    let mut heartbeat_serial: u32 = 10_000_000; // 心跳用独立序列号空间
    let mut heartbeat_send_failures = 0u32;

    loop {
        tokio::select! {
            // 接收服务端消息
            result = conn.recv() => {
                match result {
                    Ok(Some(frame)) => {
                        let proto_id = frame.header.proto_id;
                        let body = match decode_client_inbound_body(&frame, aes_key.as_ref()) {
                            Ok(body) => body,
                            Err(e) => {
                                tracing::warn!(
                                    error = %e,
                                    serial = frame.header.serial_no,
                                    proto_id,
                                    "decrypt failed, dropping inbound frame"
                                );
                                release_pending_on_inbound_decode_error(&frame, &pending);
                                continue;
                            }
                        };

                        if proto_id::is_push_proto(proto_id) {
                            // 推送消息
                            send_push_message(&push_tx, PushMessage { proto_id, body });
                        } else {
                            // 响应消息：匹配 serial number
                            let serial = frame.header.serial_no;
                            match pending.remove(&serial) {
                                Some((_, req)) => {
                                    let resp_frame = FutuFrame {
                                        header: frame.header,
                                        body,
                                    };
                                    send_response_frame(req.tx, resp_frame);
                                }
                                _ => {
                                    tracing::debug!(
                                        serial = serial,
                                        proto_id = proto_id,
                                        "unmatched response"
                                    );
                                }
                            }
                        }
                    }
                    Ok(None) => {
                        tracing::warn!("connection closed by server");
                        break;
                    }
                    Err(e) => {
                        tracing::error!(error = %e, "recv error");
                        break;
                    }
                }
            }

            // 发送命令
            cmd = cmd_rx.recv() => {
                match cmd {
                    Some(ClientCommand::Send(frame, ack)) => {
                        let result = conn.send(frame).await;
                        if let Err(e) = &result {
                            tracing::error!(error = %e, "send failed");
                        }
                        send_command_ack(ack);
                        if result.is_err() {
                            break;
                        }
                    }
                    None => {
                        tracing::info!("shutting down event loop");
                        break;
                    }
                }
            }

            // 心跳
            _ = heartbeat.tick() => {
                heartbeat_serial += 1;
                let req = futu_proto::keep_alive::Request {
                    c2s: futu_proto::keep_alive::C2s {
                        time: chrono::Utc::now().timestamp(),
                    },
                };
                let body = prost::Message::encode_to_vec(&req);
                let frame = Connection::build_frame(
                    proto_id::KEEP_ALIVE,
                    heartbeat_serial,
                    body,
                );
                let result = conn.send(frame).await;
                if let Err(e) = &result {
                    tracing::warn!(
                        error = %e,
                        consecutive_failures = heartbeat_send_failures.saturating_add(1),
                        max_failures = MAX_HEARTBEAT_SEND_FAILURES,
                        "heartbeat send failed"
                    );
                }
                if heartbeat_send_failure_should_exit(&mut heartbeat_send_failures, &result) {
                    tracing::error!(
                        consecutive_failures = heartbeat_send_failures,
                        max_failures = MAX_HEARTBEAT_SEND_FAILURES,
                        "heartbeat send failure threshold reached; closing event loop"
                    );
                    break;
                }
                if result.is_ok() {
                    tracing::trace!("heartbeat sent");
                }
            }
        }
    }

    // 清理所有 pending 请求
    pending.clear();
    tracing::info!("event loop exited");
}

/// InitConnect 握手返回的信息
#[derive(Debug, Clone)]
pub struct InitConnectInfo {
    pub server_ver: i32,
    pub login_user_id: u64,
    pub conn_id: u64,
    pub conn_aes_key: String,
    pub keep_alive_interval: i32,
}

/// 带自动重连的客户端包装
pub struct ReconnectingClient {
    config: ClientConfig,
    policy: ReconnectPolicy,
    init_ai_type: Option<i32>,
}

impl ReconnectingClient {
    pub fn new(config: ClientConfig) -> Self {
        Self {
            config,
            policy: ReconnectPolicy::default_policy(),
            init_ai_type: None,
        }
    }

    pub fn with_policy(mut self, policy: ReconnectPolicy) -> Self {
        self.policy = policy;
        self
    }

    /// Set C++ 10.7 `InitConnect.C2S.aiType` on each reconnect attempt.
    ///
    /// `None` preserves the legacy wire shape. C++ treats an absent value as
    /// `0` (non-AI); `Some(1)` marks a skills/AI caller.
    pub fn with_init_ai_type(mut self, ai_type: Option<i32>) -> Self {
        self.init_ai_type = ai_type;
        self
    }

    /// 带重连的连接循环
    ///
    /// 返回成功连接的 (FutuClient, push_rx, InitConnectInfo)。
    /// 如果达到最大重试次数则返回错误。
    pub async fn connect(&mut self) -> Result<(FutuClient, PushReceiver, InitConnectInfo)> {
        loop {
            let (mut client, push_rx) = FutuClient::new(self.config.clone());
            client = client.with_init_ai_type(self.init_ai_type);
            match client.connect().await {
                Ok(info) => {
                    self.policy.reset();
                    return Ok((client, push_rx, info));
                }
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        attempt = self.policy.attempts(),
                        "connection failed"
                    );
                    match self.policy.next_delay() {
                        Some(delay) => {
                            tracing::info!(delay_ms = delay.as_millis(), "reconnecting...");
                            tokio::time::sleep(delay).await;
                        }
                        None => {
                            return Err(FutuError::Codec(format!(
                                "max retries reached after {} attempts",
                                self.policy.attempts()
                            )));
                        }
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests;

fn hex_decode_16(hex_bytes: &[u8]) -> Option<[u8; 16]> {
    if hex_bytes.len() != 32 {
        return None;
    }
    let hex_str = std::str::from_utf8(hex_bytes).ok()?;
    let mut key = [0u8; 16];
    for i in 0..16 {
        key[i] = u8::from_str_radix(&hex_str[i * 2..i * 2 + 2], 16).ok()?;
    }
    Some(key)
}