Fix early half-close exit command finishing a stream early

orlandos-nl · Aug 30, 2024 · f265b4e · f265b4e
1 parent 350e1a4
commit f265b4e
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 12 deletions.
diff --git a/Examples/OpenSSH/Dockerfile b/Examples/OpenSSH/Dockerfile
@@ -3,7 +3,7 @@ FROM ubuntu:latest
 
 RUN apt update && apt install openssh-server sudo -y
 
-RUN echo 'ubuntu:test' | chpasswd
+RUN echo "ubuntu:test" | chpasswd
 
 RUN service ssh start
 

diff --git a/Sources/Citadel/ClientSession.swift b/Sources/Citadel/ClientSession.swift
@@ -17,7 +17,6 @@ final class ClientHandshakeHandler: ChannelInboundHandler {
     }
 
     func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
-        print(event)
         if event is UserAuthSuccessEvent {
             self.promise.succeed(())
         }

diff --git a/Sources/Citadel/TTY/Client/TTY.swift b/Sources/Citadel/TTY/Client/TTY.swift
@@ -58,9 +58,12 @@ final class ExecCommandHandler: ChannelDuplexHandler {
     typealias OutboundOut = SSHChannelData
 
     let logger: Logger
-    let onOutput: (Channel, Output) -> ()
-
-    init(logger: Logger, onOutput: @escaping (Channel, Output) -> ()) {
+    let onOutput: (Channel, Output) -> Void
+
+    init(
+        logger: Logger,
+        onOutput: @escaping (Channel, Output) -> Void
+    ) {
         self.logger = logger
         self.onOutput = onOutput
     }
@@ -160,6 +163,7 @@ extension SSHClient {
         }
 
         var hasReceivedChannelSuccess = false
+        var exitCode: Int?
 
         let handler = ExecCommandHandler(logger: logger) { channel, output in
             switch output {
@@ -168,7 +172,13 @@ extension SSHClient {
             case .stderr(let stderr):
                 streamContinuation.yield(.stderr(stderr))
             case .eof(let error):
-                streamContinuation.finish(throwing: error)
+                if let error {
+                    streamContinuation.finish(throwing: error)
+                } else if let exitCode, exitCode != 0 {
+                    streamContinuation.finish(throwing: CommandFailed(exitCode: exitCode))
+                } else {
+                    streamContinuation.finish()
+                }
             case .channelSuccess:
                 if inShell, !hasReceivedChannelSuccess {
                     let commandData = SSHChannelData(type: .channel,
@@ -177,11 +187,7 @@ extension SSHClient {
                     hasReceivedChannelSuccess = true
                 }
             case .exit(let status):
-                if status == 0 {
-                    streamContinuation.finish()
-                } else {
-                    streamContinuation.finish(throwing: CommandFailed(exitCode: status))
-                }
+                exitCode = status
             }
         }
 

diff --git a/Tests/CitadelTests/Citadel2Tests.swift b/Tests/CitadelTests/Citadel2Tests.swift
@@ -275,7 +275,8 @@ final class Citadel2Tests: XCTestCase {
             reconnect: .never
         )
 
-        _ = try await client.executeCommand("ls")
+        let output = try await client.executeCommand("ls /")
+        XCTAssertFalse(String(buffer: output).isEmpty)
 
         try await client.close()
     }