Implement an SSH Server Shell support

Package.resolved

@@ -15,8 +15,8 @@
         "repositoryURL": "https://github.com/apple/swift-atomics.git",
         "state": {
           "branch": null,
-          "revision": "ff3d2212b6b093db7f177d0855adbc4ef9c5f036",
-          "version": "1.0.3"
+          "revision": "6c89474e62719ddcc1e9614989fff2f68208fe10",
+          "version": "1.1.0"
         }
       },
       {
@@ -42,17 +42,17 @@
         "repositoryURL": "https://github.com/apple/swift-log.git",
         "state": {
           "branch": null,
-          "revision": "6fe203dc33195667ce1759bf0182975e4653ba1c",
-          "version": "1.4.4"
+          "revision": "32e8d724467f8fe623624570367e3d50c5638e46",
+          "version": "1.5.2"
         }
       },
       {
         "package": "swift-nio",
         "repositoryURL": "https://github.com/apple/swift-nio.git",
         "state": {
           "branch": null,
-          "revision": "45167b8006448c79dda4b7bd604e07a034c15c49",
-          "version": "2.48.0"
+          "revision": "2d8e6ca36fe3e8ed74b0883f593757a45463c34d",
+          "version": "2.53.0"
         }
       },
       {

Package.swift

@@ -34,6 +34,11 @@ let package = Package(
                 .product(name: "Logging", package: "swift-log"),
             ]
         ),
+        .executableTarget(
+            name: "CitadelServerExample",
+            dependencies: [
+                "Citadel"
+            ]),
         .testTarget(
             name: "CitadelTests",
             dependencies: [

Sources/Citadel/Exec/Client/ExecClient.swift

@@ -0,0 +1,124 @@
+import Foundation
+import NIO
+import NIOSSH
+
+final class TTYHandler: ChannelDuplexHandler {
+    typealias InboundIn = SSHChannelData
+    typealias InboundOut = ByteBuffer
+    typealias OutboundIn = ByteBuffer
+    typealias OutboundOut = SSHChannelData
+
+    let maxResponseSize: Int
+    var isIgnoringInput = false
+    var response = ByteBuffer()
+    let done: EventLoopPromise<ByteBuffer>
+
+    init(
+        maxResponseSize: Int,
+        done: EventLoopPromise<ByteBuffer>
+    ) {
+        self.maxResponseSize = maxResponseSize
+        self.done = done
+    }
+
+    func handlerAdded(context: ChannelHandlerContext) {
+        context.channel.setOption(ChannelOptions.allowRemoteHalfClosure, value: true).whenFailure { error in
+            context.fireErrorCaught(error)
+        }
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        switch event {
+        case is SSHChannelRequestEvent.ExitStatus:
+            ()
+        default:
+            context.fireUserInboundEventTriggered(event)
+        }
+    }
+
+    func handlerRemoved(context: ChannelHandlerContext) {
+        done.succeed(response)
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let data = self.unwrapInboundIn(data)
+
+        guard case .byteBuffer(var bytes) = data.data, !isIgnoringInput else {
+            return
+        }
+
+        switch data.type {
+        case .channel:
+            if
+                response.readableBytes + bytes.readableBytes > maxResponseSize
+            {
+                isIgnoringInput = true
+                done.fail(CitadelError.commandOutputTooLarge)
+                return
+            }
+
+            // Channel data is forwarded on, the pipe channel will handle it.
+            response.writeBuffer(&bytes)
+            return
+        case .stdErr:
+            done.fail(TTYSTDError(message: bytes))
+        default:
+            ()
+        }
+    }
+
+    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
+        let data = self.unwrapOutboundIn(data)
+        context.write(self.wrapOutboundOut(SSHChannelData(type: .channel, data: .byteBuffer(data))), promise: promise)
+    }
+}
+
+extension SSHClient {
+    /// Executes a command on the remote server. This will return the output of the command. If the command fails, the error will be thrown. If the output is too large, the command will fail.
+    /// - Parameters:
+    ///  - command: The command to execute.
+    /// - maxResponseSize: The maximum size of the response. If the response is larger, the command will fail.
+    public func executeCommand(_ command: String, maxResponseSize: Int = .max) async throws -> ByteBuffer {
+        let promise = eventLoop.makePromise(of: ByteBuffer.self)
+
+        let channel: Channel
+
+        do {
+            channel = try await eventLoop.flatSubmit {
+                let createChannel = self.eventLoop.makePromise(of: Channel.self)
+                self.session.sshHandler.createChannel(createChannel) { channel, _ in
+                    channel.pipeline.addHandlers(
+                        TTYHandler(
+                            maxResponseSize: maxResponseSize,
+                            done: promise
+                        )
+                    )
+                }
+
+                self.eventLoop.scheduleTask(in: .seconds(15)) {
+                    createChannel.fail(CitadelError.channelCreationFailed)
+                }
+
+                return createChannel.futureResult
+            }.get()
+        } catch {
+            promise.fail(error)
+            throw error
+        }
+
+        // We need to exec a thing.
+        let execRequest = SSHChannelRequestEvent.ExecRequest(
+            command: command,
+            wantReply: true
+        )
+
+        return try await eventLoop.flatSubmit {
+            channel.triggerUserOutboundEvent(execRequest).whenFailure { [channel] error in
+                channel.close(promise: nil)
+                promise.fail(error)
+            }
+
+            return promise.futureResult
+        }.get()
+    }
+}

Sources/Citadel/SFTP/Server/SFTPServer.swift

@@ -69,18 +69,18 @@ public protocol SFTPDelegate {
     func rename(oldPath: String, newPath: String, flags: UInt32, context: SSHContext) async throws -> SFTPStatusCode
 }
 
-struct SFTPServerSubsystem {
+enum SFTPServerSubsystem {
     static func setupChannelHanders(
         channel: Channel,
-        delegate: SFTPDelegate,
+        sftp: SFTPDelegate,
         logger: Logger,
         username: String?
     ) -> EventLoopFuture<Void> {
         let deserializeHandler = ByteToMessageHandler(SFTPMessageParser())
         let serializeHandler = MessageToByteHandler(SFTPMessageSerializer())
         let sftpInboundHandler = SFTPServerInboundHandler(
             logger: logger,
-            delegate: delegate,
+            delegate: sftp,
             eventLoop: channel.eventLoop,
             username: username
         )

Sources/Citadel/Server.swift

@@ -22,12 +22,14 @@ final class SubsystemHandler: ChannelDuplexHandler {
     typealias OutboundIn = SSHChannelData
     typealias OutboundOut = SSHChannelData
 
+    let shell: ShellDelegate?
     let sftp: SFTPDelegate?
     let eventLoop: EventLoop
     var configured: EventLoopPromise<Void>
 
-    init(sftp: SFTPDelegate?, eventLoop: EventLoop) {
+    init(sftp: SFTPDelegate?, shell: ShellDelegate?, eventLoop: EventLoop) {
         self.sftp = sftp
+        self.shell = shell
         self.eventLoop = eventLoop
         self.configured = eventLoop.makePromise()
     }
@@ -38,24 +40,55 @@ final class SubsystemHandler: ChannelDuplexHandler {
         }
     }
 
+    func handlerRemoved(context: ChannelHandlerContext) {
+        self.configured.succeed(())
+    }
+
     func channelInactive(context: ChannelHandlerContext) {
         context.fireChannelInactive()
     }
 
     func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
         switch event {
+        case let event as SSHChannelRequestEvent.ExecRequest:
+            context.fireUserInboundEventTriggered(event)
+        case is SSHChannelRequestEvent.ShellRequest:
+            guard let shell = shell, let parent = context.channel.parent else {
+                _ = context.channel.triggerUserOutboundEvent(ChannelFailureEvent()).flatMap {
+                    self.configured.succeed(())
+                    return context.channel.close()
+                }
+                return
+            }
+
+            parent.pipeline.handler(type: NIOSSHHandler.self).flatMap { handler in
+                ShellServerSubsystem.setupChannelHanders(
+                    channel: context.channel,
+                    shell: shell,
+                    logger: .init(label: "nl.orlandos.citadel.sftp-server"),
+                    username: handler.username
+                )
+            }.flatMap { () -> EventLoopFuture<Void> in
+                let promise = context.eventLoop.makePromise(of: Void.self)
+                context.channel.triggerUserOutboundEvent(ChannelSuccessEvent(), promise: promise)
+                self.configured.succeed(())
+                return promise.futureResult
+            }.whenFailure { _ in
+                context.channel.triggerUserOutboundEvent(ChannelFailureEvent(), promise: nil)
+            }
         case let event as SSHChannelRequestEvent.SubsystemRequest:
             switch event.subsystem {
             case "sftp":
                 guard let sftp = sftp, let parent = context.channel.parent else {
                     context.channel.close(promise: nil)
+                    self.configured.succeed(())
                     return
                 }
 
                 parent.pipeline.handler(type: NIOSSHHandler.self).flatMap { handler in
                     SFTPServerSubsystem.setupChannelHanders(
                         channel: context.channel,
-                        delegate: sftp,
+                        sftp: sftp,
                         logger: .init(label: "nl.orlandos.citadel.sftp-server"),
                         username: handler.username
                     )
@@ -93,6 +126,7 @@ final class SubsystemHandler: ChannelDuplexHandler {
 final class CitadelServerDelegate {
     var sftp: SFTPDelegate?
     var exec: ExecDelegate?
+    var shell: ShellDelegate?
 
     fileprivate init() {}
 
@@ -103,6 +137,7 @@ final class CitadelServerDelegate {
 
             handlers.append(SubsystemHandler(
                 sftp: self.sftp,
+                shell: self.shell,
                 eventLoop: channel.eventLoop
             ))
 
@@ -149,6 +184,10 @@ public final class SSHServer {
         self.delegate.exec = delegate
     }
 
+    public func enableShell(withDelegate delegate: ShellDelegate) {
+        self.delegate.shell = delegate
+    }
+
     /// Closes the SSH Server, stopping new connections from coming in.
     public func close() async throws {
         try await channel.close()

Sources/Citadel/Shell/Server/ShellDelegate.swift

@@ -0,0 +1,101 @@
+import NIO
+import NIOSSH
+import Logging
+
+public enum ShellClientEvent {
+    case stdin(ByteBuffer)
+}
+
+public enum ShellServerEvent {
+    case stdout(ByteBuffer)
+}
+
+public protocol ShellDelegate {
+    func startShell(
+        reading stream: AsyncStream<ShellClientEvent>,
+        context: SSHContext
+    ) async throws -> AsyncThrowingStream<ShellServerEvent, Error>
+}
+
+fileprivate struct ShellContinuation {
+    var continuation: AsyncStream<ShellClientEvent>.Continuation!
+}
+
+final class ShellServerInboundHandler: ChannelInboundHandler {
+    typealias InboundIn = ByteBuffer
+
+    let logger: Logger
+    let delegate: ShellDelegate
+    let username: String?
+    let eventLoop: EventLoop
+    fileprivate var streamWriter: ShellContinuation
+    let stream: AsyncStream<ShellClientEvent>
+
+    init(logger: Logger, delegate: ShellDelegate, eventLoop: EventLoop, username: String?) {
+        self.logger = logger
+        self.delegate = delegate
+        self.username = username
+        self.eventLoop = eventLoop
+
+        var streamWriter = ShellContinuation()
+        self.stream = AsyncStream { continuation in
+            streamWriter.continuation = continuation
+        }
+        self.streamWriter = streamWriter
+    }
+
+    func handlerAdded(context: ChannelHandlerContext) {
+        let channel = context.channel
+
+        let done = context.eventLoop.makePromise(of: Void.self)
+        done.completeWithTask {
+            let output = try await self.delegate.startShell(
+                reading: self.stream,
+                context: SSHContext(username: self.username)
+            )
+
+            for try await chunk in output {
+                switch chunk {
+                case .stdout(let data):
+                    try await channel.writeAndFlush(data)
+                }
+            }
+        }
+
+        done.futureResult.whenFailure(context.fireErrorCaught)
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        streamWriter.continuation.yield(.stdin(unwrapInboundIn(data)))
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        switch event {
+        default:
+            context.fireUserInboundEventTriggered(event)
+        }
+    }
+}
+
+enum ShellServerSubsystem {
+    static func setupChannelHanders(
+        channel: Channel,
+        shell: ShellDelegate,
+        logger: Logger,
+        username: String?
+    ) -> EventLoopFuture<Void> {
+        let shellInboundHandler = ShellServerInboundHandler(
+            logger: logger,
+            delegate: shell,
+            eventLoop: channel.eventLoop,
+            username: username
+        )
+
+        return channel.pipeline.addHandlers(
+            SSHChannelDataUnwrapper(),
+            SSHOutboundChannelDataWrapper(),
+            shellInboundHandler,
+            CloseErrorHandler(logger: logger)
+        )
+    }
+}