From d90abcfb50c2121795539f65fa51a9705e6bba48 Mon Sep 17 00:00:00 2001
From: Eugene Pankov <e@ajenti.org>
Date: Sun, 20 Nov 2022 13:59:29 +0100
Subject: [PATCH] fixed missing CHANNEL_CLOSE messages - #459

---
 Cargo.lock                                      |  4 ++--
 warpgate-protocol-ssh/Cargo.toml                |  4 ++--
 .../src/client/channel_session.rs               | 17 +++++++++++++++--
 warpgate-protocol-ssh/src/client/mod.rs         |  9 +++------
 4 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2b8ac25..4d94671 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3199,9 +3199,9 @@ dependencies = [
 
 [[package]]
 name = "russh"
-version = "0.35.0-beta.5"
+version = "0.35.0-beta.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "355d7ad01489febf7f475495a898b49987f4fa79b2a70bfee3e109194a5ed93b"
+checksum = "0c130a5e8cb48fd94d117612590cd029924e3d968403c4064c709795fac0d52f"
 dependencies = [
  "aes",
  "aes-gcm",
diff --git a/warpgate-protocol-ssh/Cargo.toml b/warpgate-protocol-ssh/Cargo.toml
index 0b45fc0..29d0e58 100644
--- a/warpgate-protocol-ssh/Cargo.toml
+++ b/warpgate-protocol-ssh/Cargo.toml
@@ -12,8 +12,8 @@ bimap = "0.6"
 bytes = "1.2"
 dialoguer = "0.10"
 futures = "0.3"
-russh = { version = "0.35.0-beta.5", features = ["vendored-openssl"] }
-# russh = { version = "0.35.0-beta.5", features = ["vendored-openssl"], path = "../../russh/russh"}
+russh = { version = "0.35.0-beta.6", features = ["vendored-openssl"] }
+# russh = { version = "0.35.0-beta.6", features = ["vendored-openssl"], path = "../../russh/russh"}
 russh-keys = { version = "0.23.0-beta.1", features = ["vendored-openssl"] }
 # russh-keys = { version = "0.23.0-beta.1", features = ["vendored-openssl"], path = "../../russh/russh-keys" }
 sea-orm = { version = "0.10.3", features = [
diff --git a/warpgate-protocol-ssh/src/client/channel_session.rs b/warpgate-protocol-ssh/src/client/channel_session.rs
index b1c04aa..458d47c 100644
--- a/warpgate-protocol-ssh/src/client/channel_session.rs
+++ b/warpgate-protocol-ssh/src/client/channel_session.rs
@@ -16,6 +16,7 @@ pub struct SessionChannel {
     ops_rx: UnboundedReceiver<ChannelOperation>,
     events_tx: UnboundedSender<RCEvent>,
     session_id: SessionId,
+    closed: bool,
 }
 
 impl SessionChannel {
@@ -32,6 +33,7 @@ impl SessionChannel {
             ops_rx,
             events_tx,
             session_id,
+            closed: false,
         }
     }
 
@@ -110,7 +112,6 @@ impl SessionChannel {
                             )).map_err(|_| SshClientError::MpscError)?;
                         }
                         Some(russh::ChannelMsg::Close) => {
-                            self.events_tx.send(RCEvent::Close(self.channel_id)).map_err(|_| SshClientError::MpscError)?;
                             break;
                         },
                         Some(russh::ChannelMsg::Success) => {
@@ -147,19 +148,31 @@ impl SessionChannel {
                             warn!("unhandled channel message: {:?}", msg);
                         }
                         None => {
-                            self.events_tx.send(RCEvent::Close(self.channel_id)).map_err(|_| SshClientError::MpscError)?;
                             break
                         },
                     }
                 }
             }
         }
+        self.close()?;
+        Ok(())
+    }
+
+    fn close(&mut self) -> Result<(), SshClientError> {
+        if !self.closed {
+            let _ = self
+                .events_tx
+                .send(RCEvent::Close(self.channel_id))
+                .map_err(|_| SshClientError::MpscError);
+            self.closed = true;
+        }
         Ok(())
     }
 }
 
 impl Drop for SessionChannel {
     fn drop(&mut self) {
+        let _ = self.close();
         info!(channel=%self.channel_id, session=%self.session_id, "Closed");
     }
 }
diff --git a/warpgate-protocol-ssh/src/client/mod.rs b/warpgate-protocol-ssh/src/client/mod.rs
index a9bd83f..1a4ccb9 100644
--- a/warpgate-protocol-ssh/src/client/mod.rs
+++ b/warpgate-protocol-ssh/src/client/mod.rs
@@ -241,15 +241,12 @@ impl RemoteClient {
             op => {
                 let mut channel_pipes = self.channel_pipes.lock().await;
                 match channel_pipes.get(&channel_id) {
-                    Some(tx) => match tx.send(op) {
-                        Ok(_) => {}
-                        Err(SendError(_)) => {
+                    Some(tx) => {
+                        if tx.send(op).is_err() {
                             channel_pipes.remove(&channel_id);
                         }
-                    },
-                    None => {
-                        debug!(channel=%channel_id, "operation for unknown channel")
                     }
+                    None => debug!(channel=%channel_id, "operation for unknown channel"),
                 }
             }
         }