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;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::trace;

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

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

pub struct AbstractWsMode {}

impl AbstractWsMode {
    //创建链接
    pub async fn ws_connect_async<T, PI, PO>(bool_v1: Arc<AtomicBool>,
                                             address_url: String,
                                             lable: String,
                                             subscribe_array: Vec<String>,
                                             mut write_rx: UnboundedReceiver<Message>,
                                             read_tx: UnboundedSender<ResponseData>,
                                             message_text: T,
                                             message_ping: PI,
                                             message_pong: PO) -> Result<(), Error>
        where T: Fn(String) -> Option<ResponseData> + Copy,
              PI: Fn(Vec<u8>) -> Option<ResponseData> + Copy,
              PO: Fn(Vec<u8>) -> Option<ResponseData> + Copy
    {
        //1.是否走代理
        /*******走代理：根据环境变量配置来决定,如果配置了走代理，没有配置不走*******/
        let proxy = match proxy::ParsingDetail::env_proxy(ProxyEnum::WS) {
            ProxyResponseEnum::NO => {
                // trace!("非 代理");
                None
            }
            ProxyResponseEnum::YES(proxy) => {
                // trace!("代理");
                Option::from(proxy)
            }
        };

        loop {
            let (ws_stream, _) = connect_async(address_url.clone(), proxy).await?;
            trace!("WebSocket 握手完成。");
            let (mut write, mut read) = ws_stream.split();

            //订阅写入(包括订阅信息 )
            trace!("订阅内容：{:?}",subscribe_array.clone());
            for s in &subscribe_array {
                write.send(Message::Text(s.parse().unwrap())).await?;
            }

            //将socket 的写操作与 写通道链接起来，将数据以ok的结构体封装进行传递
            // let stdin_to_ws = write_rx.map(Ok).forward(write);
            // Writing task
            let stdin_to_ws = async {
                while let Some(message) = write_rx.next().await {
                    write.send(message).await?;
                }
                Ok::<(), tokio_tungstenite::tungstenite::Error>(())
            };
            let ws_to_stdout = async {
                while let Some(message) = read.next().await {
                    let response_data = AbstractWsMode::analysis_message(message, message_text, message_ping, message_pong);
                    // let response_data = func(message);
                    if response_data.is_some() {
                        let mut data = response_data.unwrap();
                        data.label = lable.clone();
                        let code = data.code.clone();
                        if code.as_str() == "-1" {} else if code.as_str() == "200" {
                            if bool_v1.load(Ordering::Relaxed) {
                                read_tx.unbounded_send(data).unwrap();
                            }
                        }
                    }
                }
                Ok::<(), tokio_tungstenite::tungstenite::Error>(())
            };
            // let ws_to_stdout = {
            //     trace!("---1");
            //     //读，循环读取，然后拿到 message,，然后开启异步处理 message，
            //     let result = read.for_each(|message| async {
            //         let response_data = func(message);
            //         if response_data.is_some() {
            //             let mut data = response_data.unwrap();
            //             data.label = lable.clone();
            //             let code = data.code.clone();
            //             if code.as_str() == "-1" {
            //                 // let close_frame = CloseFrame {
            //                 //     code: CloseCode::Normal,
            //                 //     reason: Cow::Borrowed("Bye bye"),
            //                 // };
            //                 // let close_message = Message::Close(Some(close_frame));
            //                 // write.send(close_message);
            //             } else if code.as_str() == "200" {
            //                 read_tx.unbounded_send(data).unwrap();
            //             }
            //         }
            //     });
            //     trace!("---3");
            //     result
            // };

            //必须操作。，因为不同于其他的高级语言，有自动内存管理，所以为了防范地址改变，所以需要做此处理
            pin_mut!(stdin_to_ws, ws_to_stdout,);
            future::select(stdin_to_ws, ws_to_stdout).await;
            trace!("---5");
            trace!("---4");
            trace!("重启...");
        }
        // return Ok(());
    }
    //心跳包
    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"))
                    }
                }
            ).expect("发送失败");
            trace!("发送指令-心跳：{:?}",h_type);
        }
    }
    //数据解析
    pub fn analysis_message<T, PI, PO>(message: Result<Message, Error>,
                                       message_text: T,
                                       message_ping: PI,
                                       message_pong: PO) -> Option<ResponseData>
        where T: Fn(String) -> Option<ResponseData>,
              PI: Fn(Vec<u8>) -> Option<ResponseData>,
              PO: 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)) => {
                //二进制WebSocket消息
                let message_str = format!("Binary:{:?}", s);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(),
                                               "2".to_string(),
                                               message_str, "".to_string()))
            }
            Ok(Message::Close(c)) => {
                let message_str = format!("关闭指令:{:?}", c);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(),
                                               "0".to_string(),
                                               message_str, "".to_string()))
            }
            Ok(Message::Frame(f)) => {
                //原始帧 正常读取数据不会读取到该 信息类型
                let message_str = format!("意外读取到原始帧:{:?}", f);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(),
                                               "-2".to_string(),
                                               message_str, "".to_string()))
            }
            Err(e) => {
                let message_str = format!("服务器响应:{:?}", e);
                trace!("{:?}",message_str);
                Option::from(ResponseData::new("".to_string(),
                                               "-1".to_string(),
                                               message_str, "".to_string()))
            }
        }
    }
    //发送数据
    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: futures_channel::mpsc::UnboundedSender<Message>) {
    let _str = serde_json::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
}