Implement --echo argument to enable the echoing of shell io on the target session

Author: TimeToogo

tunshell-client/src/client.rs

@@ -1,5 +1,5 @@
 #[cfg(not(target_arch = "wasm32"))]
-use crate::p2p;
+use crate::{p2p, ShellServerConfig};
 use crate::{
     AesStream, ClientMode, Config, HostShell, RelayStream, ServerStream, ShellKey, TunnelStream,
 };
@@ -312,10 +312,12 @@ impl Client {
 
     #[cfg(not(target_arch = "wasm32"))]
     async fn start_shell_server(&self, peer_socket: Box<dyn TunnelStream>) -> Result<u8> {
-        crate::ShellServer::new()?
-            .run(peer_socket, ShellKey::new(self.config.encryption_key()))
-            .await
-            .and_then(|_| Ok(0))
+        crate::ShellServer::new(ShellServerConfig {
+            echo_stdout: self.config.echo_stdout(),
+        })?
+        .run(peer_socket, ShellKey::new(self.config.encryption_key()))
+        .await
+        .and_then(|_| Ok(0))
     }
 
     #[cfg(target_arch = "wasm32")]
@@ -362,6 +364,7 @@ mod tests {
             443,
             "test",
             true,
+            false,
         );
         let mut client = Client::new(config, HostShell::new().unwrap());
 

tunshell-client/src/config.rs

@@ -13,6 +13,7 @@ pub struct Config {
     relay_ws_port: u16,
     server_connection_timeout: Duration,
     direct_connection_timeout: Duration,
+    echo_stdout: bool,
     enable_direct_connection: bool,
     dangerous_disable_relay_server_verification: bool,
 }
@@ -44,6 +45,15 @@ impl Config {
             .unwrap_or("443".to_owned())
             .parse::<u16>()
             .expect("could not parse arg (7) as WebSocket port");
+        let mut echo_stdout = false;
+
+        // Parse optional args
+        while let Some(arg) = args.next() {
+            match arg.as_str() {
+                "--echo" => echo_stdout = true,
+                arg @ _ => panic!(format!("Unknown argument: {}", arg)),
+            }
+        }
 
         if session_key.len() < 10 {
             panic!("session key is too short")
@@ -62,6 +72,7 @@ impl Config {
             encryption_key,
             server_connection_timeout: Duration::from_millis(DEFAULT_SERVER_CONNECT_TIMEOUT),
             direct_connection_timeout: Duration::from_millis(DEFAULT_DIRECT_CONNECT_TIMEOUT),
+            echo_stdout,
             enable_direct_connection: true,
             dangerous_disable_relay_server_verification: false,
         }
@@ -75,6 +86,7 @@ impl Config {
         relay_ws_port: u16,
         encryption_key: &str,
         enable_direct_connection: bool,
+        echo_stdout: bool,
     ) -> Self {
         Self {
             mode,
@@ -86,6 +98,7 @@ impl Config {
             server_connection_timeout: Duration::from_millis(DEFAULT_SERVER_CONNECT_TIMEOUT),
             direct_connection_timeout: Duration::from_millis(DEFAULT_DIRECT_CONNECT_TIMEOUT),
             enable_direct_connection,
+            echo_stdout,
             dangerous_disable_relay_server_verification: false,
         }
     }
@@ -126,6 +139,10 @@ impl Config {
         self.direct_connection_timeout
     }
 
+    pub fn echo_stdout(&self) -> bool {
+        self.echo_stdout
+    }
+
     pub fn enable_direct_connection(&self) -> bool {
         self.enable_direct_connection
     }

tunshell-client/src/server/mod.rs

@@ -30,6 +30,7 @@ mod tests {
             443,
             "test",
             true,
+            false
         );
 
         let result = Runtime::new()

tunshell-client/src/shell/server/mod.rs

@@ -4,7 +4,7 @@ use anyhow::{Error, Result};
 use futures::stream::StreamExt;
 use log::*;
 use std::time::Duration;
-use tokio::time;
+use tokio::{io::AsyncWriteExt, time};
 use tokio_util::compat::*;
 
 mod fallback;
@@ -23,11 +23,17 @@ use pty::*;
 
 type ShellStream = ShellServerStream<Compat<Box<dyn TunnelStream>>>;
 
-pub(crate) struct ShellServer {}
+pub(crate) struct ShellServer {
+    conf: ShellServerConfig,
+}
+
+pub(crate) struct ShellServerConfig {
+    pub(crate) echo_stdout: bool,
+}
 
 impl ShellServer {
-    pub(crate) fn new() -> Result<ShellServer> {
-        Ok(ShellServer {})
+    pub(crate) fn new(conf: ShellServerConfig) -> Result<ShellServer> {
+        Ok(ShellServer { conf })
     }
 
     pub(crate) async fn run(self, stream: Box<dyn TunnelStream>, key: ShellKey) -> Result<()> {
@@ -108,6 +114,11 @@ impl ShellServer {
         mut shell: Box<dyn Shell + Send + 'a>,
     ) -> Result<()> {
         let mut buff = [0u8; 1024];
+        let mut host_stdout = if self.conf.echo_stdout {
+            Some(tokio::io::stdout())
+        } else {
+            None
+        };
 
         loop {
             info!("waiting for shell message");
@@ -122,6 +133,12 @@ impl ShellServer {
                     },
                     Ok(read) => {
                         info!("read {} bytes from stdout", read);
+
+                        if let Some(host_stdout) = host_stdout.as_mut() {
+                            host_stdout.write_all(&buff[..read]).await?;
+                            host_stdout.flush().await?;
+                        }
+
                         stream.write(&ShellServerMessage::Stdout(buff[..read].to_vec())).await?;
                         info!("sent {} bytes to client shell", read);
                     },
@@ -167,9 +184,13 @@ mod tests {
     use tokio::time::timeout;
     use tunshell_shared::Message;
 
+    fn new_shell_server() -> ShellServer {
+        ShellServer::new(ShellServerConfig { echo_stdout: true }).unwrap()
+    }
+
     #[test]
     fn test_new_shell_server() {
-        ShellServer::new().unwrap();
+        new_shell_server();
     }
 
     #[test]
@@ -186,8 +207,7 @@ mod tests {
             );
 
             let mock_stream = Cursor::new(mock_data).compat();
-            ShellServer::new()
-                .unwrap()
+            new_shell_server()
                 .run(Box::new(mock_stream), ShellKey::new("MyKey"))
                 .await
                 .expect_err("client key should be rejected");
@@ -203,9 +223,7 @@ mod tests {
 
             timeout(
                 Duration::from_millis(5000),
-                ShellServer::new()
-                    .unwrap()
-                    .run(Box::new(mock_stream), ShellKey::new("CorrectKey")),
+                new_shell_server().run(Box::new(mock_stream), ShellKey::new("CorrectKey")),
             )
             .await
             .unwrap()
@@ -230,9 +248,7 @@ mod tests {
 
             timeout(
                 Duration::from_millis(5000),
-                ShellServer::new()
-                    .unwrap()
-                    .run(Box::new(mock_stream), ShellKey::new("CorrectKey")),
+                new_shell_server().run(Box::new(mock_stream), ShellKey::new("CorrectKey")),
             )
             .await
             .unwrap()
@@ -289,7 +305,7 @@ mod tests {
             );
 
             let mock_stream = Cursor::new(mock_data).compat();
-            let server = ShellServer::new().unwrap();
+            let server = new_shell_server();
 
             server
                 .run(Box::new(mock_stream), ShellKey::new("CorrectKey"))
@@ -331,7 +347,7 @@ mod tests {
             );
 
             let mock_stream = Cursor::new(mock_data).compat();
-            let server = ShellServer::new().unwrap();
+            let server = new_shell_server();
 
             server
                 .run(Box::new(mock_stream), ShellKey::new("CorrectKey"))

tunshell-client/src/wasm.rs

@@ -93,6 +93,7 @@ pub async fn tunshell_init_client(config: BrowserConfig) {
         5000,
         443,
         &config.encryption_key,
+        false,
         false
     );
 

tunshell-tests/tests/valid_direct_connection.rs

@@ -44,6 +44,7 @@ fn test() {
             config.api_port,
             "mock_encryption_key",
             true,
+            false,
         );
         target_config.set_dangerous_disable_relay_server_verification(true);
         let mut target = client::Client::new(target_config, target_shell.clone());
@@ -58,6 +59,7 @@ fn test() {
             config.api_port,
             "mock_encryption_key",
             true,
+            false,
         );
         local_config.set_dangerous_disable_relay_server_verification(true);
         let mut local = client::Client::new(local_config, local_shell.clone());

tunshell-tests/tests/valid_relay_connection.rs

@@ -44,6 +44,7 @@ fn test() {
             config.api_port,
             "mock_encryption_key",
             false,
+            false,
         );
         target_config.set_dangerous_disable_relay_server_verification(true);
         let mut target = client::Client::new(target_config, target_shell.clone());
@@ -58,6 +59,7 @@ fn test() {
             config.api_port,
             "mock_encryption_key",
             false,
+            false,
         );
         local_config.set_dangerous_disable_relay_server_verification(true);
         let mut local = client::Client::new(local_config, local_shell.clone());