Forward ChannelMsg::Close to channel before dropping sender

When the remote side sends CHANNEL_CLOSE, the client and server
handlers were removing the channel entry (dropping the mpsc Sender)
without first forwarding ChannelMsg::Close. Consumers on
Channel::wait() would see None with no preceding Close message,
unlike CHANNEL_EOF which correctly sends ChannelMsg::Eof first.

Send ChannelMsg::Close through the channel's mpsc before removing
it, matching the existing CHANNEL_EOF pattern.

Author: coreyleavitt

russh/src/channels/mod.rs

@@ -153,6 +153,9 @@ impl std::fmt::Debug for ChannelReadHalf {
 
 impl ChannelReadHalf {
     /// Awaits an incoming [`ChannelMsg`], this method returns [`None`] if the channel has been closed.
+    ///
+    /// When the remote end sends `CHANNEL_CLOSE`, a [`ChannelMsg::Close`] is
+    /// delivered before this method returns [`None`].
     pub async fn wait(&mut self) -> Option<ChannelMsg> {
         self.receiver.recv().await
     }
@@ -581,6 +584,9 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> Channel<S> {
     }
 
     /// Awaits an incoming [`ChannelMsg`], this method returns [`None`] if the channel has been closed.
+    ///
+    /// When the remote end sends `CHANNEL_CLOSE`, a [`ChannelMsg::Close`] is
+    /// delivered before this method returns [`None`].
     pub async fn wait(&mut self) -> Option<ChannelMsg> {
         self.read_half.wait().await
     }

russh/src/client/encrypted.rs

@@ -365,6 +365,12 @@ impl Session {
                     // will not be released.
                     enc.close(channel_num)?;
                 }
+                // Forward the close to the channel before removing it, so that
+                // consumers waiting on `Channel::wait()` receive an explicit
+                // `ChannelMsg::Close` instead of just seeing `None`.
+                if let Some(chan) = self.channels.get(&channel_num) {
+                    let _ = chan.send(ChannelMsg::Close).await;
+                }
                 self.channels.remove(&channel_num);
                 client.channel_close(channel_num, self).await
             }

russh/src/server/encrypted.rs

@@ -602,6 +602,12 @@ impl Session {
                 if let Some(ref mut enc) = self.common.encrypted {
                     enc.channels.remove(&channel_num);
                 }
+                // Forward the close to the channel before removing it, so that
+                // consumers waiting on `Channel::wait()` receive an explicit
+                // `ChannelMsg::Close` instead of just seeing `None`.
+                if let Some(chan) = self.channels.get(&channel_num) {
+                    chan.send(ChannelMsg::Close).await.unwrap_or(())
+                }
                 self.channels.remove(&channel_num);
                 debug!("handler.channel_close {channel_num:?}");
                 handler.channel_close(channel_num, self).await

russh/src/tests.rs

@@ -455,6 +455,13 @@ mod channels {
                     panic!("Unexpected message {msg:?}");
                 }
 
+                // After the server closes the channel, we should receive an
+                // explicit Close message before the channel stream ends.
+                let msg = ch.wait().await.unwrap();
+                assert!(
+                    matches!(msg, ChannelMsg::Close),
+                    "expected Close, got {msg:?}"
+                );
                 assert!(ch.wait().await.is_none());
                 c
             },
@@ -463,6 +470,78 @@ mod channels {
         .await;
     }
 
+    /// Verify that the server-side CHANNEL_CLOSE handler delivers
+    /// `ChannelMsg::Close` before the channel stream ends.
+    #[tokio::test]
+    async fn test_server_receives_close_on_client_close() {
+        #[derive(Debug)]
+        struct Client {}
+
+        impl client::Handler for Client {
+            type Error = crate::Error;
+
+            async fn check_server_key(
+                &mut self,
+                _server_public_key: &crate::keys::ssh_key::PublicKey,
+            ) -> Result<bool, Self::Error> {
+                Ok(true)
+            }
+        }
+
+        struct ServerHandle {
+            channel: Option<tokio::sync::oneshot::Sender<Channel<server::Msg>>>,
+        }
+
+        impl server::Handler for ServerHandle {
+            type Error = crate::Error;
+
+            async fn auth_publickey(
+                &mut self,
+                _: &str,
+                _: &crate::keys::ssh_key::PublicKey,
+            ) -> Result<server::Auth, Self::Error> {
+                Ok(server::Auth::Accept)
+            }
+
+            async fn channel_open_session(
+                &mut self,
+                channel: Channel<server::Msg>,
+                _session: &mut server::Session,
+            ) -> Result<bool, Self::Error> {
+                if let Some(tx) = self.channel.take() {
+                    tx.send(channel).unwrap();
+                }
+                Ok(true)
+            }
+        }
+
+        let (tx, rx) = tokio::sync::oneshot::channel::<Channel<server::Msg>>();
+        let sh = ServerHandle { channel: Some(tx) };
+
+        test_session(
+            Client {},
+            sh,
+            |c| async move {
+                let ch = c.channel_open_session().await.unwrap();
+                ch.close().await.unwrap();
+                c
+            },
+            |s| async move {
+                let mut ch = rx.await.unwrap();
+                // The server should receive an explicit Close message
+                // when the client closes the channel.
+                let msg = ch.wait().await.unwrap();
+                assert!(
+                    matches!(msg, ChannelMsg::Close),
+                    "expected Close, got {msg:?}"
+                );
+                assert!(ch.wait().await.is_none());
+                s
+            },
+        )
+        .await;
+    }
+
     #[tokio::test]
     async fn test_channel_window_size() {
         #[derive(Debug)]

russh/tests/test_backpressure.rs

@@ -135,6 +135,7 @@ impl russh::server::Handler for Server {
             while let Ok(_) = rx.changed().await {
                 match channel.wait().await {
                     Some(ChannelMsg::Data { .. }) => (),
+                    Some(ChannelMsg::Close) | None => break,
                     other => panic!("unexpected message {other:?}"),
                 }
             }

russh/tests/test_data_stream.rs

@@ -72,7 +72,7 @@ impl ChannelDataCopy for ChannelHalves {
                     match msg {
                         ChannelMsg::WindowAdjusted { .. } => {}
                         ChannelMsg::Data { data } => buf.extend(&*data),
-                        ChannelMsg::Eof => break,
+                        ChannelMsg::Eof | ChannelMsg::Close => break,
                         msg => panic!("Got unexpected message: {msg:?}"),
                     }
                 }