1 files changed, 22 insertions, 6 deletions
diff --git a/src/main.rs b/src/main.rs
index cc4163d..d93b24f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -42,7 +42,7 @@ async fn start_client(url: Url) -> Result<()> {
         match tokio_tungstenite::connect_async(&url).await {
             Ok((stream, _resp)) => {
                 eprintln!("connected to server");
-                webcat("client >>>", stream, &mut stdin_receiver).await?;
+                webcat(Mode::Client, stream, &mut stdin_receiver).await?;
             }
             Err(err) => {
                 eprintln!("failed to connect to server: {}", err);
@@ -58,6 +58,8 @@ async fn start_client(url: Url) -> Result<()> {
     }
 }
 
+const SERVER_ACCEPTED_NEW_CONN: &str = "accepted new client";
+
 async fn start_server(port: u16) -> Result<()> {
     let listen_addr: SocketAddr = ([127, 0, 0, 1], port).into();
     let listener = TcpListener::bind(listen_addr).await?;
@@ -70,10 +72,10 @@ async fn start_server(port: u16) -> Result<()> {
 
     loop {
         let (stream, _client_addr) = listener.accept().await?;
-        eprintln!("accepted TCP connection");
         match tokio_tungstenite::accept_async(stream).await {
             Ok(stream) => {
-                webcat("server <<<", stream, &mut stdin_receiver).await?;
+                println!("{}", SERVER_ACCEPTED_NEW_CONN);
+                webcat(Mode::Server, stream, &mut stdin_receiver).await?;
             }
             Err(err) => {
                 eprintln!("failed to upgrade {}", err);
@@ -89,8 +91,14 @@ async fn send_stdin_to_channel(sender: UnboundedSender<String>) -> Result<()> {
     Ok(())
 }
 
+#[derive(PartialEq)]
+enum Mode {
+    Server,
+    Client,
+}
+
 async fn webcat<S: Stream<Item = Result<Message, WsError>> + Sink<Message>>(
-    name: &str,
+    mode: Mode,
     stream: S,
     receiver: &mut UnboundedReceiver<String>,
 ) -> Result<()> {
@@ -98,7 +106,12 @@ async fn webcat<S: Stream<Item = Result<Message, WsError>> + Sink<Message>>(
     loop {
         tokio::select! {
             Some(line) = receiver.recv() => {
-                if write.send(Message::Text(line)).await.is_err() {
+                if atty::isnt(atty::Stream::Stdin) && mode == Mode::Client && line == SERVER_ACCEPTED_NEW_CONN {
+                    // Assuming the MITM setup described in the README, we
+                    // disconnect and reconnect the client to prevent the
+                    // application client and server from desynchronizing.
+                    return Ok(());
+                } else if write.send(Message::Text(line)).await.is_err() {
                     eprintln!("failed to send message");
                 }
             }
@@ -111,7 +124,10 @@ async fn webcat<S: Stream<Item = Result<Message, WsError>> + Sink<Message>>(
                                 if !atty::is(atty::Stream::Stdout) {
                                     // if stdout is redirected we also print the
                                     // messages to stderr for convenience
-                                    eprintln!("{} {}", name, text);
+                                    eprintln!("{} {}", match mode {
+                                        Mode::Server => "server <<<",
+                                        Mode::Client => "client >>>",
+                                    }, text);
                                 }
                             }
                             other => {