as-rust/exchanges/src/socket_tool.rs

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration};

use chrono::Utc;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use futures_util::{future, pin_mut, SinkExt, StreamExt};
use futures_util::stream::{SplitSink, SplitStream};
use ring::hmac;
use serde_json::{json, Value};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio::time::Instant;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::{error, info, trace};

use crate::proxy;
use crate::proxy::{ProxyEnum, ProxyResponseEnum};
use crate::response_base::ResponseData;

#[derive(Debug)]
pub enum HeartbeatType {
    Ping,
    Pong,
    Custom(String),
}

pub struct AbstractWsMode {}

impl AbstractWsMode {
    pub async fn ws_connected<T, PI, PO, F, B, Future>(write_to_socket_rx_arc: Arc<Mutex<UnboundedReceiver<Message>>>,
                                                    is_first_login: bool,
                                                    label: String,
                                                    is_shutdown_arc: Arc<AtomicBool>,
                                                    handle_function: &F,
                                                    subscribe_array: Vec<String>,
                                                    ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
                                                    message_text: T,
                                                    message_ping: PI,
                                                    message_pong: PO,
                                                    message_binary: B)
        where T: Fn(String) -> Option<ResponseData> + Copy,
              PI: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              PO: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              F: Fn(ResponseData) -> Future + Clone,
              B: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              Future: future::Future<Output=()> + Send + 'static,
    {
        let (ws_write, mut ws_read) = ws_stream.split();
        let ws_write_arc = Arc::new(Mutex::new(ws_write));

        // 将socket 的写操作与【写通道（外部向socket写）】链接起来，将数据以ok的结构体封装进行传递
        // 这里是形成链式操作，如果要将外界的信息传进来（使用socket查单、下单之类的，部分交易所可以支持），就要这要弄
        let mut write_to_socket_rx = write_to_socket_rx_arc.lock().await;
        let ws_write_channel_clone = Arc::clone(&ws_write_arc);
        let stdin_to_ws = async {
            while let Some(message) = write_to_socket_rx.next().await {
                let mut write_lock2 = ws_write_channel_clone.lock().await;
                write_lock2.send(message).await?;
            }
            Ok::<(), Error>(())
        };
        // 如果不需要事先登录，则直接订阅消息
        if !is_first_login {
            info!("不需要登录，订阅内容：{:?}", subscribe_array.clone());
            for s in &subscribe_array {
                let mut write_lock = ws_write_arc.lock().await;
                write_lock.send(Message::Text(s.parse().unwrap())).await.expect("订阅消息失败");
            }
        }

        let ws_write_inner = Arc::clone(&ws_write_arc);
        let ws_to_stdout = async {
            while let Some(message) = ws_read.next().await {
                if !is_shutdown_arc.load(Ordering::Relaxed) {
                    continue;
                }

                let response_data = AbstractWsMode::analysis_message(message, message_text, message_ping, message_pong, message_binary);
                // let response_data = func(message);
                if response_data.is_some() {
                    let mut data = response_data.unwrap();
                    data.label = label.clone();

                    let code = data.code.clone();

                    if code == 200 {
                        let mut data_c = data.clone();
                        data_c.ins = Instant::now();
                        data_c.time = Utc::now().timestamp_millis();

                        handle_function(data_c).await;
                    }

                    /*
                        200 -正确返回
                       -200 -登录成功
                       -201 -订阅成功
                       -300 -客户端收到服务器心跳ping，需要响应
                       -301 -客户端收到服务器心跳pong，需要响应
                       -302 -客户端收到服务器心跳自定义，需要响应自定义
                    */
                    match code {
                        200 => {
                            let mut data_c = data.clone();
                            data_c.ins = Instant::now();
                            data_c.time = Utc::now().timestamp_millis();

                            handle_function(data_c).await;
                        }
                        -200 => {
                            //登录成功
                            info!("ws登录成功：{:?}", data);
                            info!("订阅内容：{:?}", subscribe_array.clone());
                            for s in &subscribe_array {
                                let mut write_lock = ws_write_arc.lock().await;
                                write_lock.send(Message::Text(s.parse().unwrap())).await.expect("订阅消息失败");
                            }
                            info!("订阅完成！");
                        }
                        -201 => {
                            //订阅成功
                            trace!("订阅成功：{:?}", data);
                        }
                        -300 => {
                            //服务器发送心跳 ping 给客户端，客户端需要pong回应
                            trace!("服务器响应-ping");
                            if data.data != Value::Null {
                                let mut ws_write = ws_write_inner.lock().await;
                                ws_write.send(Message::Pong(Vec::from(data.data.to_string()))).await?;
                                trace!("客户端回应服务器-pong");
                            }
                        }
                        -301 => {
                            //服务器发送心跳 pong 给客户端，客户端需要ping回应
                            trace!("服务器响应-pong");
                            if data.data != Value::Null {
                                let mut ws_write = ws_write_inner.lock().await;
                                ws_write.send(Message::Ping(Vec::from(data.data.to_string()))).await?;
                                trace!("客户端回应服务器-ping");
                            }
                        }
                        -302 => {
                            //客户端收到服务器心跳自定义，需要响应自定义
                            trace!("特定字符心跳，特殊响应：{:?}", data);
                            let mut ws_write = ws_write_inner.lock().await;
                            ws_write.send(Message::Text(data.data.to_string())).await?;
                            trace!("特殊字符心跳-回应完成");
                        }
                        _ => {
                            trace!("未知：{:?}",data);
                        }
                    }
                }
            }
            Ok::<(), Error>(())
        };

        //必须操作。，因为不同于其他的高级语言，有自动内存管理，所以为了防范地址改变，所以需要做此处理
        pin_mut!(stdin_to_ws, ws_to_stdout,);
        future::select(stdin_to_ws, ws_to_stdout).await;
    }

    //创建链接
    pub async fn ws_connect_async<T, PI, PO, F, B, Future>(is_shutdown_arc: Arc<AtomicBool>,
                                                        handle_function: F,
                                                        address_url: String,
                                                        is_first_login: bool,
                                                        label: String,
                                                        subscribe_array: Vec<String>,
                                                        write_to_socket_rx_arc: Arc<Mutex<UnboundedReceiver<Message>>>,
                                                        message_text: T,
                                                        message_ping: PI,
                                                        message_pong: PO,
                                                        message_binary: B)
        where T: Fn(String) -> Option<ResponseData> + Copy,
              PI: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              PO: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              B: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              F: Fn(ResponseData) -> Future + Clone,
              Future: future::Future<Output=()> + Send + 'static,
    {
        //1.是否走代理
        /*******走代理：根据环境变量配置来决定,如果配置了走代理，没有配置不走*******/
        let proxy = match proxy::ParsingDetail::env_proxy(ProxyEnum::WS) {
            ProxyResponseEnum::NO => {
                // trace!("非 代理");
                None
            }
            ProxyResponseEnum::YES(proxy) => {
                // trace!("代理");
                Option::from(proxy)
            }
        };

        match connect_async(address_url.clone(), proxy).await {
            Ok((ws_stream, _)) => {
                info!("socket 链接成功，{}。", address_url);

                Self::ws_connected(write_to_socket_rx_arc,
                                   is_first_login,
                                   label,
                                   is_shutdown_arc,
                                   &handle_function,
                                   subscribe_array.clone(),
                                   ws_stream,
                                   message_text,
                                   message_ping,
                                   message_pong,
                                   message_binary).await;
            }
            Err(e) => {
                error!("WebSocket 握手失败：{:?}", e);
            }
        }
    }

    //心跳包
    pub async fn ping_or_pong(write_tx_clone: Arc<Mutex<UnboundedSender<Message>>>, h_type: HeartbeatType, millis: u64) {
        loop {
            tokio::time::sleep(Duration::from_millis(millis)).await;
            let write_tx_clone = write_tx_clone.lock().await;
            write_tx_clone.unbounded_send(
                match h_type {
                    HeartbeatType::Ping => {
                        Message::Ping(Vec::from("Ping"))
                    }
                    HeartbeatType::Pong => {
                        Message::Pong(Vec::from("Pong"))
                    }
                    HeartbeatType::Custom(ref str) => {
                        Message::Text(str.parse().unwrap())
                    }
                }
            ).expect("发送失败");
            trace!("发送指令-心跳：{:?}",h_type);
        }
    }
    //数据解析
    pub fn analysis_message<T, PI, PO, B>(message: Result<Message, Error>,
                                       message_text: T,
                                       message_ping: PI,
                                       message_pong: PO,
                                       message_binary: B) -> Option<ResponseData>
        where T: Fn(String) -> Option<ResponseData>,
              PI: Fn(Vec<u8>) -> Option<ResponseData>,
              PO: Fn(Vec<u8>) -> Option<ResponseData>,
              B: Fn(Vec<u8>) -> Option<ResponseData>
    {
        match message {
            Ok(Message::Text(text)) => message_text(text),
            Ok(Message::Ping(pi)) => message_ping(pi),
            Ok(Message::Pong(po)) => message_pong(po),
            Ok(Message::Binary(s)) => message_binary(s), //二进制WebSocket消息
            Ok(Message::Close(c)) => {
                let message_str = format!("关闭指令:{:?}", c);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(), 0, message_str, Value::Null))
            }
            Ok(Message::Frame(f)) => {
                //原始帧 正常读取数据不会读取到该 信息类型
                let message_str = format!("意外读取到原始帧:{:?}", f);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(), -2, message_str, Value::Null))
            }
            Err(e) => {
                let message_str = format!("服务器响应:{:?}", e);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(), -1, message_str, Value::Null))
            }
        }
    }
    //发送数据
    pub async fn send_subscribe(write_tx_clone: Arc<Mutex<UnboundedSender<Message>>>, message: Message) -> bool {
        let write_tx_clone = write_tx_clone.lock().await;
        write_tx_clone.unbounded_send(message.clone()).unwrap();
        trace!("发送指令：{:?}",message);
        true
    }
}

//创建链接
pub async fn ws_connect_async(address_url: String) -> (SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
                                                       SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>) {
    //1.是否走代理
    /*******走代理：根据环境变量配置来决定,如果配置了走代理，没有配置不走*******/
    let proxy = match proxy::ParsingDetail::env_proxy(ProxyEnum::WS) {
        ProxyResponseEnum::NO => {
            trace!("非 代理");
            None
        }
        ProxyResponseEnum::YES(proxy) => {
            trace!("代理");
            Option::from(proxy)
        }
    };

    let (ws_stream, _) = connect_async(address_url, proxy).await.expect("链接失败！");
    trace!("WebSocket 握手完成。");
    ws_stream.split()
}


pub async fn client(add_url: String) {
    let proxy = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(
        127,
        0,
        0,
        1)
    ), 7890);


    //创建通道 开启线程，向通道写入数据
    let (write_tx, write_rx) = futures_channel::mpsc::unbounded();
    let (read_tx, read_rx) = futures_channel::mpsc::unbounded();
    tokio::spawn(write_sell(write_tx));


    //创建socket,，并且读写分离
    let (ws_stream, _) = connect_async(add_url, Option::from(proxy)).await.expect("Failed to connect");
    trace!("WebSocket handshake has been successfully completed");
    let (write, read) = ws_stream.split();

    //将socket 的写操作与 写通道链接起来，将数据以ok的结构体封装进行传递
    let stdin_to_ws = write_rx.map(Ok).forward(write);
    let ws_to_stdout = {
        trace!("---1");
        //读，循环读取，然后拿到 message,，然后开启异步处理 message，
        let result = read.for_each(|message| async {
            read_tx.unbounded_send(message.unwrap()).unwrap();
        });
        trace!("---3");
        result
    };

    tokio::spawn(read_sell(read_rx));

    //必须操作。，因为不同于其他的高级语言，有自动内存管理，所以为了防范地址改变，所以需要做此处理
    pin_mut!(stdin_to_ws, ws_to_stdout);
    future::select(stdin_to_ws, ws_to_stdout).await;
}


//模拟 业务场景中 发送指令给目标交易所
async fn write_sell(tx: UnboundedSender<Message>) {
    let _str = json!({
                "op": "subscribe",
                "args": [
                        {
                        // "channel":"orders",
                        // "instType":"SWAP",
                        // "instFamily":"BTC-USDT"
                        "channel":"books5",
                        "instId":"BTC-USDT"
                        }
                    ]
            });
    let str_array: Vec<String> = vec![
        // log_in_to_str(),
        // str.to_string(),
    ];

    let i = 0;
    loop {
        if str_array.len() > i {
            let send_str = str_array.get(i).unwrap();
            tx.unbounded_send(Message::Text(send_str.to_string())).unwrap();
        }
        tokio::time::sleep(Duration::from_secs(5)).await;
        tx.unbounded_send(Message::Ping(Vec::from("Ping"))).unwrap();
        tx.unbounded_send(Message::Ping(Vec::from("Pong"))).unwrap();
    }
}


async fn read_sell(mut rx: futures_channel::mpsc::UnboundedReceiver<Message>) {
    loop {
        if let Some(message) = rx.next().await {
            match message {
                Message::Text(s) => {
                    trace!("Text: {}", s);
                }
                Message::Binary(s) => {
                    trace!("Binary: {:?}", s);
                }
                Message::Ping(s) => {
                    trace!("Ping: {:?}", s);
                }
                Message::Pong(s) => {
                    trace!("Pong: {:?}", s);
                }
                Message::Close(s) => {
                    trace!("Close: {:?}", s);
                }
                Message::Frame(s) => {
                    trace!("Frame: {:?}", s);
                }
            }
        }
        tokio::time::sleep(Duration::from_millis(1)).await
    }
}


pub fn log_in_to_str() -> String {
    let mut login_json_str = "".to_string();

    let access_key: String = "".to_string();
    let secret_key: String = "".to_string();
    let passphrase: String = "".to_string();

    if access_key.len() > 0 || secret_key.len() > 0 || passphrase.len() > 0 {
        let timestamp = Utc::now().timestamp().to_string();
        // 时间戳 + 请求类型+ 请求参数字符串
        let message = format!("{}GET{}", timestamp, "/users/self/verify");
        let hmac_key = ring::hmac::Key::new(hmac::HMAC_SHA256, secret_key.as_bytes());
        let result = ring::hmac::sign(&hmac_key, &message.as_bytes());
        let sign = base64::encode(result);

        let login_json = json!({
                              "op": "login",
                                "args": [{
                                "apiKey": access_key,
                                "passphrase": passphrase,
                                "timestamp": timestamp,
                                "sign": sign  }]
                        });

        // trace!("---login_json：{0}", login_json.to_string());
        // trace!("--登录：{}", login_json.to_string());
        login_json_str = login_json.to_string();
    }
    login_json_str
}