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;

/// 客户端配置
#[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,
}

/// 请求上下文（等待响应的 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: mpsc::UnboundedSender<PushMessage>,
    cmd_tx: Option<mpsc::Sender<ClientCommand>>,
    conn_aes_key: parking_lot::Mutex<Option<[u8; 16]>>,
    conn_id: AtomicU64,
}

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

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

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);
    }
}

impl FutuClient {
    /// 创建客户端（未连接状态）
    ///
    /// 返回 (client, push_rx)，其中 push_rx 用于接收服务端推送。
    pub fn new(config: ClientConfig) -> (Self, mpsc::UnboundedReceiver<PushMessage>) {
        let (push_tx, push_rx) = mpsc::unbounded_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),
        };

        (client, push_rx)
    }

    /// 连接到 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 = futu_proto::init_connect::Request {
            c2s: futu_proto::init_connect::C2s {
                client_ver: 100,
                client_id: self.config.client_id.clone(),
                recv_notify: Some(self.config.recv_notify),
                packet_enc_algo: Some(0),
                push_proto_fmt: Some(0),
                programming_language: Some(String::new()),
            },
        };
        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");
            encrypt::rsa_private_decrypt(rsa_key, &resp_frame.body)?
        } else {
            resp_frame.body.to_vec()
        };

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

        // 检查返回码
        let ret_type = resp.ret_type;
        if ret_type != 0 {
            return Err(FutuError::ServerError {
                ret_type,
                msg: resp.ret_msg.unwrap_or_default(),
            });
        }

        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(),
                key_hex = hex_str(key_bytes),
                "received AES key"
            );
            if key_bytes.len() == 16 {
                // key 直接是 16 字节原始值
                let mut key = [0u8; 16];
                key.copy_from_slice(key_bytes);
                *self.conn_aes_key.lock() = Some(key);
            } else if key_bytes.len() == 32 {
                // key 是 32 字符十六进制编码的 16 字节
                if let Some(key) = hex_decode_16(key_bytes) {
                    *self.conn_aes_key.lock() = Some(key);
                } else {
                    tracing::warn!(
                        "AES key is 32 chars but not valid hex, using raw first 16 bytes"
                    );
                    let mut key = [0u8; 16];
                    key.copy_from_slice(&key_bytes[..16]);
                    *self.conn_aes_key.lock() = Some(key);
                }
            } else {
                tracing::warn!(
                    key_len = key_bytes.len(),
                    "unexpected AES key length, using raw bytes (truncated/padded to 16)"
                );
                let mut key = [0u8; 16];
                let copy_len = key_bytes.len().min(16);
                key[..copy_len].copy_from_slice(&key_bytes[..copy_len]);
                *self.conn_aes_key.lock() = Some(key);
            }
        }

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

        // 启动后台任务：心跳、消息接收
        let keep_alive_interval = Duration::from_secs(info.keep_alive_interval as u64);
        self.start_background_tasks(conn, keep_alive_interval);

        Ok(info)
    }

    /// 发送请求并等待响应
    pub async fn request(&self, proto_id: u32, body: Vec<u8>) -> Result<FutuFrame> {
        let serial = self.next_serial();

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

        let mut frame = FutuFrame::new(proto_id, serial, Bytes::from(final_body));
        // 使用明文的 SHA1
        frame.header.body_sha1 = 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();
            cmd_tx
                .send(ClientCommand::Send(frame, ack_tx))
                .await
                .map_err(|_| FutuError::NotInitialized)?;
            ack_rx
                .await
                .map_err(|_| FutuError::Codec("send ack failed".into()))?;
        } else {
            self.pending.remove(&serial);
            return Err(FutuError::NotInitialized);
        }

        // 等待响应（带超时）
        let resp = tokio::time::timeout(Duration::from_secs(12), resp_rx)
            .await
            .map_err(|_| {
                self.pending.remove(&serial);
                FutuError::Timeout
            })?
            .map_err(|_| FutuError::Codec("response channel closed".into()))?;

        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: mpsc::UnboundedSender<PushMessage>,
    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; // 心跳用独立序列号空间

    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) {
                            // 推送消息
                            let _ = push_tx.send(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,
                                };
                                let _ = req.tx.send(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");
                        }
                        let _ = ack.send(());
                        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,
                );
                if let Err(e) = conn.send(frame).await {
                    tracing::error!(error = %e, "heartbeat send failed");
                    break;
                }
                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,
}

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

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

    /// 带重连的连接循环
    ///
    /// 返回成功连接的 (FutuClient, push_rx, InitConnectInfo)。
    /// 如果达到最大重试次数则返回错误。
    pub async fn connect(
        &mut self,
    ) -> Result<(
        FutuClient,
        mpsc::UnboundedReceiver<PushMessage>,
        InitConnectInfo,
    )> {
        loop {
            let (mut client, push_rx) = FutuClient::new(self.config.clone());
            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_str(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

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)
}