more logic in cli connect

src/cli.rs

@@ -1,6 +1,8 @@
 use crate::client::*;
 use hbb_common::{
     config::PeerConfig,
+    config::READ_TIMEOUT,
+    futures::{SinkExt, StreamExt},
     log,
     message_proto::*,
     protobuf::Message as _,
@@ -87,23 +89,38 @@ impl Interface for Session {
 }
 
 #[tokio::main(flavor = "current_thread")]
-pub async fn connect_test(
+pub async fn connect_test(id: &str, key: String, token: String) {
-    id: &str,
-    key: String,
-    token: String,
-) {
     let (sender, mut receiver) = mpsc::unbounded_channel::<Data>();
     let handler = Session::new(&id, sender);
-    if let Err(err) = crate::client::Client::start(
+    match crate::client::Client::start(id, &key, &token, ConnType::PORT_FORWARD, handler).await {
-        id,
+        Err(err) => {
-        &key,
-        &token,
-        ConnType::PORT_FORWARD,
-        handler,
-    ).await {
             log::error!("Failed to connect {}: {}", &id, err);
-    } else {
+        }
+        Ok((mut stream, direct)) => {
+            log::info!("direct: {}", direct);
             // rpassword::prompt_password("Input anything to exit").ok();
+            loop {
+                tokio::select! {
+                    res = hbb_common::timeout(READ_TIMEOUT, stream.next()) => match res {
+                        Err(_) => {
+                            log::error!("Timeout");
+                            break;
+                        }
+                        Ok(Some(Ok(bytes))) => {
+                            let msg_in = Message::parse_from_bytes(&bytes).unwrap();
+                            match msg_in.union {
+                                Some(message::Union::Hash(hash)) => {
+                                    log::info!("Got hash");
+                                    break;
+                                }
+                                _ => {}
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            }
+        }
     }
 }