From be8a401f88df3e9ffdcbe0df322d40b2add6bf1a Mon Sep 17 00:00:00 2001
From: Ben Hagen <ben@ottomatic.io>
Date: Sat, 14 Dec 2024 17:53:15 +0100
Subject: [PATCH] fix: Close channel correctly when receiving
 SSH_Message_Channel_Close

Fixes #116

When receiving a SSH_Message_Channel_Close message, the channel was not
being properly closed. This caused the stdout stream to remain open,
preventing the drain() operation from completing.

The fix:
1. Let the channel handle the close message properly
2. Added test to verify channel close behavior
---
 lib/src/ssh_client.dart                |  8 ++++++-
 test/src/channel/ssh_channel_test.dart | 30 ++++++++++++++++++++++++++
 test/test_utils.dart                   | 10 +++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 test/src/channel/ssh_channel_test.dart

diff --git a/lib/src/ssh_client.dart b/lib/src/ssh_client.dart
index 9ddc183..700240f 100644
--- a/lib/src/ssh_client.dart
+++ b/lib/src/ssh_client.dart
@@ -21,6 +21,7 @@ import 'package:dartssh2/src/message/msg_userauth.dart';
 import 'package:dartssh2/src/ssh_message.dart';
 import 'package:dartssh2/src/socket/ssh_socket.dart';
 import 'package:dartssh2/src/ssh_userauth.dart';
+import 'package:meta/meta.dart';
 
 /// https://datatracker.ietf.org/doc/html/rfc4252#section-8
 typedef SSHPasswordRequestHandler = FutureOr<String?> Function();
@@ -486,6 +487,10 @@ class SSHClient {
     }
   }
 
+  /// Handles a raw SSH packet. This method is only exposed for testing purposes.
+  @visibleForTesting
+  void handlePacket(Uint8List packet) => _handlePacket(packet);
+
   void _sendMessage(SSHMessage message) {
     printTrace?.call('-> $socket: $message');
     _transport.sendPacket(message.encode());
@@ -792,7 +797,8 @@ class SSHClient {
     printTrace?.call('<- $socket: $message');
     final channel = _channels[message.recipientChannel];
     if (channel != null) {
-      channel.close();
+      channel.handleMessage(message);
+      //channel.close();
       _channels.remove(message.recipientChannel);
       _channelIdAllocator.release(message.recipientChannel);
     }
diff --git a/test/src/channel/ssh_channel_test.dart b/test/src/channel/ssh_channel_test.dart
new file mode 100644
index 0000000..7fcc41c
--- /dev/null
+++ b/test/src/channel/ssh_channel_test.dart
@@ -0,0 +1,30 @@
+import 'package:dartssh2/dartssh2.dart';
+import 'package:test/test.dart';
+
+import '../../test_utils.dart';
+
+void main() {
+  group('SSHChannel', () {
+    late SSHClient client;
+    late SSHSession session;
+
+    setUp(() async {
+      client = await getTestClient();
+      await client.authenticated;
+      session = await client.shell();
+    });
+
+    tearDown(() {
+      client.close();
+    });
+
+    test('stdout stream handles remote channel close correctly', () async {
+      final drainFuture = session.stdout.drain<void>();
+
+      final closeMessage = createChannelCloseMessage(0);
+      client.handlePacket(closeMessage);
+
+      await drainFuture;
+    }, timeout: const Timeout(Duration(seconds: 1)));
+  });
+}
diff --git a/test/test_utils.dart b/test/test_utils.dart
index 0550184..f4c96c0 100644
--- a/test/test_utils.dart
+++ b/test/test_utils.dart
@@ -1,6 +1,8 @@
 import 'dart:io';
+import 'dart:typed_data';
 
 import 'package:dartssh2/dartssh2.dart';
+import 'package:dartssh2/src/message/msg_channel.dart';
 
 /// A honeypot that accepts all passwords and public-keys
 Future<SSHClient> getHoneypotClient() async {
@@ -40,3 +42,11 @@ Future<List<SSHKeyPair>> getTestKeyPairs() async {
 String fixture(String path) {
   return File('test/fixtures/$path').readAsStringSync();
 }
+
+/// Create a [SSH_Message_Channel_Close] message.
+Uint8List createChannelCloseMessage(int recipientChannel) {
+  final message = SSH_Message_Channel_Close(
+    recipientChannel: recipientChannel,
+  );
+  return message.encode();
+}