Merge pull request #54 from orlandos-nl/jo/direct-tcpip-server

Add DirectTCPIP support to SSH servers

Author: Joannis

.gitignore

@@ -5,3 +5,4 @@
 /*.xcodeproj
 xcuserdata/
 Package.resolved
+.vscode/launch.json

Package.resolved

@@ -1,17 +1,19 @@
 import NIO
 import NIOSSH
+import Logging
 
 final class ClientHandshakeHandler: ChannelInboundHandler {
     typealias InboundIn = Any
 
     private let promise: EventLoopPromise<Void>
+    let logger = Logger(label: "nl.orlandos.citadel.handshake")
 
     /// A future that will be fulfilled when the handshake is complete.
     public var authenticated: EventLoopFuture<Void> {
         promise.futureResult
     }
 
-    init(eventLoop: EventLoop) {
+    init(eventLoop: EventLoop, loginTimeout: TimeAmount) {
         let promise = eventLoop.makePromise(of: Void.self)
         self.promise = promise
     }
@@ -54,7 +56,10 @@ final class SSHClientSession {
         algorithms: SSHAlgorithms = SSHAlgorithms(),
         protocolOptions: Set<SSHProtocolOption> = []
     ) async throws -> SSHClientSession {
-        let handshakeHandler = ClientHandshakeHandler(eventLoop: channel.eventLoop)
+        let handshakeHandler = ClientHandshakeHandler(
+            eventLoop: channel.eventLoop,
+            loginTimeout: .seconds(10)
+        )
         var clientConfiguration = SSHClientConfiguration(
             userAuthDelegate: authenticationMethod(),
             serverAuthDelegate: hostKeyValidator
@@ -101,7 +106,10 @@ final class SSHClientSession {
         group: EventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1),
         connectTimeout: TimeAmount = .seconds(30)
     ) async throws -> SSHClientSession {
-        let handshakeHandler = ClientHandshakeHandler(eventLoop: group.next())
+        let handshakeHandler = ClientHandshakeHandler(
+            eventLoop: group.next(),
+            loginTimeout: .seconds(10)
+        )
         var clientConfiguration = SSHClientConfiguration(
             userAuthDelegate: authenticationMethod(),
             serverAuthDelegate: hostKeyValidator

Sources/Citadel/ClientSession.swift

@@ -0,0 +1,62 @@
+import NIO
+import NIOSSH
+
+fileprivate final class ProxyChannelHandler: ChannelOutboundHandler {
+    typealias OutboundIn = ByteBuffer
+
+    private let write: (ByteBuffer, EventLoopPromise<Void>?) -> Void
+
+    init(write: @escaping (ByteBuffer, EventLoopPromise<Void>?) -> Void) {
+        self.write = write
+    }
+
+    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
+        let data = self.unwrapOutboundIn(data)
+        write(data, promise)
+    }
+}
+
+public protocol DirectTCPIPDelegate {
+    func initializeDirectTCPIPChannel(_ channel: Channel, request: SSHChannelType.DirectTCPIP, context: SSHContext) -> EventLoopFuture<Void>
+}
+
+public struct DirectTCPIPForwardingDelegate: DirectTCPIPDelegate {
+    internal enum Error: Swift.Error {
+        case forbidden
+    }
+
+    public var whitelistedHosts: [String]?
+    public var whitelistedPorts: [Int]?
+
+    public init() {}
+
+    public func initializeDirectTCPIPChannel(_ channel: Channel, request: SSHChannelType.DirectTCPIP, context: SSHContext) -> EventLoopFuture<Void> {
+        if let whitelistedHosts, !whitelistedHosts.contains(request.targetHost) {
+            return channel.eventLoop.makeFailedFuture(Error.forbidden)
+        }
+
+        if let whitelistedPorts, !whitelistedPorts.contains(request.targetPort) {
+            return channel.eventLoop.makeFailedFuture(Error.forbidden)
+        }
+
+        return ClientBootstrap(group: channel.eventLoop)
+            .connect(host: request.targetHost, port: request.targetPort)
+            .flatMap { remote in
+                channel.pipeline.addHandlers([
+                    DataToBufferCodec()
+                ]).flatMap {
+                    channel.pipeline.addHandler(ProxyChannelHandler { data, promise in
+                        remote.writeAndFlush(data, promise: promise)
+                    })
+                }.flatMap {
+                    remote.pipeline.addHandler(ProxyChannelHandler { [weak channel] data, promise in
+                        guard let channel else {
+                            promise?.fail(ChannelError.ioOnClosedChannel)
+                            return
+                        }
+                        channel.writeAndFlush(data, promise: promise)
+                    })
+                }
+            }
+    }
+}

Sources/Citadel/DirectTCPIP.swift …ectTCPIP/Client/DirectTCPIP+Client.swiftSources/Citadel/DirectTCPIP.swift renamed to Sources/Citadel/DirectTCPIP/Client/DirectTCPIP+Client.swift Sources/Citadel/DirectTCPIP.swift renamed to Sources/Citadel/DirectTCPIP/Client/DirectTCPIP+Client.swift

@@ -93,6 +93,7 @@ final class SubsystemHandler: ChannelDuplexHandler {
 final class CitadelServerDelegate {
     var sftp: SFTPDelegate?
     var exec: ExecDelegate?
+    var directTCPIP: DirectTCPIPDelegate?
 
     fileprivate init() {}
 
@@ -109,7 +110,19 @@ final class CitadelServerDelegate {
             handlers.append(ExecHandler(delegate: exec, username: username))
 
             return channel.pipeline.addHandlers(handlers)
-        case .directTCPIP, .forwardedTCPIP:
+        case .directTCPIP(let request):
+            guard let delegate = directTCPIP else {
+                return channel.eventLoop.makeFailedFuture(CitadelError.unsupported)
+            }
+
+            return channel.pipeline.addHandler(DataToBufferCodec()).flatMap {
+                return delegate.initializeDirectTCPIPChannel(
+                    channel,
+                    request: request,
+                    context: SSHContext(username: username)
+                )
+            }
+        case .forwardedTCPIP:
             return channel.eventLoop.makeFailedFuture(CitadelError.unsupported)
         }
     }
@@ -148,6 +161,10 @@ public final class SSHServer {
     public func enableExec(withDelegate delegate: ExecDelegate) {
         self.delegate.exec = delegate
     }
+
+    public func enableDirectTCPIP(withDelegate delegate: DirectTCPIPDelegate) {
+        self.delegate.directTCPIP = delegate
+    }
 
     /// Closes the SSH Server, stopping new connections from coming in.
     public func close() async throws {