feat: first commit

src/main.rs

@@ -0,0 +1,94 @@
+#[macro_use]
+extern crate rust_util;
+
+use clap::{App, Arg};
+use tokio_tungstenite::connect_async;
+use rust_util::XResult;
+use futures_util::{StreamExt, SinkExt};
+use tokio_tungstenite::tungstenite::Message;
+use tokio::sync::mpsc::channel;
+
+enum SimpleMessage {
+    End,
+    Msg(Message),
+}
+
+#[tokio::main]
+async fn main() {
+    let matches = App::new("")
+        .arg(Arg::with_name("connect").short("c").long("connect").takes_value(true).help("WebSocket URL"))
+        .get_matches();
+
+    let connect = matches.value_of("connect").unwrap_or_else(|| {
+        failure_and_exit!("Arg --connect is required.");
+    });
+
+    let r = connect_to_and_loop(connect).await;
+    match r {
+        Ok(_) => success!("Success"),
+        Err(e) => failure!("Error: {}", e),
+    }
+}
+
+async fn connect_to_and_loop(url: &str) -> XResult<()> {
+    let (ws_stream, _) = connect_async(url).await?;
+    let (mut ws_write, mut ws_read) = ws_stream.split();
+
+    let (sender, mut receiver) = channel(2);
+
+    let cloned_sender = sender.clone();
+    let end_simple_msg_loop = || async {
+        cloned_sender.send(SimpleMessage::End).await.ok();
+    };
+
+    tokio::spawn(async move {
+        loop {
+            let a = receiver.recv().await;
+            match a {
+                None => {}
+                Some(simple_msg) => {
+                    match simple_msg {
+                        SimpleMessage::End => {
+                            break;
+                        }
+                        SimpleMessage::Msg(msg) => {
+                            ws_write.send(msg).await.unwrap();
+                        }
+                    }
+                }
+            }
+        }
+    });
+
+    loop {
+        let next = ws_read.next().await;
+        match next {
+            None => {
+                information!("WebSocket ended.");
+                end_simple_msg_loop().await;
+                return Ok(());
+            }
+            Some(Err(e)) => {
+                end_simple_msg_loop().await;
+                return simple_error!("WebSocket got error: {}", e);
+            }
+            Some(Ok(message)) => {
+                match message {
+                    Message::Ping(msg) => {
+                        information!("Received ping message: {:?}", msg);
+                        let r = sender.send(SimpleMessage::Msg(Message::Pong(msg))).await;
+                        r.ok();
+                    }
+                    Message::Pong(msg) => {
+                        information!("Received pong message: {:?}", msg);
+                    }
+                    Message::Close(close_frame) => {}
+                    Message::Binary(msg) => {}
+                    Message::Text(msg) => {}
+                }
+            }
+        }
+    }
+
+    Ok(())
+}