Skip to content

Commit

Permalink
SSH Config extra flags, including DynamicForward
Browse files Browse the repository at this point in the history 
- Added pool support for DynamicForward tunnels
  • Loading branch information
Carlos Cabanero committed Feb 1, 2022
1 parent a083ff6 commit 5a18d59
Show file tree
Hide file tree
Showing 6 changed files with 349 additions and 120 deletions.
52 changes: 17 additions & 35 deletions Blink/Commands/ssh/SSHConfig.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,8 +136,8 @@ struct SSHCommand: ParsableCommand {
"Forward stdio to the specified destination",
valueName: "host:port"
),
transform: { try StdioForwardInfo($0) })
var stdioHostAndPort: StdioForwardInfo?
transform: { try BindAddressInfo($0) })
var stdioHostAndPort: BindAddressInfo?

@Option(
name: [.customShort("o", allowingJoined: true)],
Expand Down Expand Up @@ -184,7 +184,7 @@ struct SSHCommand: ParsableCommand {
valueName: "port"
)
)
var dynamicForwardingPort: UInt16?
var dynamicForward: [String] = []

// Identity
@Option(
Expand Down Expand Up @@ -299,9 +299,23 @@ extension SSHCommand {
params["remoteforward"] = self.remoteForward
}

if !self.dynamicForward.isEmpty {
params["dynamicforward"] = dynamicForward
}

if agentForward {
params["forwardagent"] = "yes"
}

if !command.isEmpty {
params["remotecommand"] = command.joined(separator: " ")
}

if disableTTY {
params["requesttty"] = "no"
} else if forceTTY {
params["requesttty"] = "force"
}

return try BKSSHHost(content: params)
}
Expand All @@ -326,38 +340,6 @@ extension SSHCommand {
}
}

struct StdioForwardInfo: Equatable {
let remotePort: UInt16
let bindAddress: String

private let pattern = #"^(?<bindAddress>\[([\w:.]+)\]|([\w.]+)):(?<remotePort>\d+)$"#

init(_ info: String) throws {
let regex = try! NSRegularExpression(pattern: pattern)

guard let match = regex.firstMatch(in: info,
range: NSRange(location: 0, length: info.count))
else {
throw ValidationError("Missing <bind_address>:<remoteport> for stdio forwarding.")
}

guard let r = Range(match.range(withName: "remotePort"), in: info),
let remotePort = UInt16(info[r])
else {
throw ValidationError("Invalid remote port.")
}
self.remotePort = remotePort

guard let r = Range(match.range(withName: "bindAddress"), in: info)
else {
throw ValidationError("Invalid bind address.")
}
var bindAddress = String(info[r])
bindAddress.removeAll(where: { $0 == "[" || $0 == "]" })
self.bindAddress = bindAddress
}
}

enum SSHControlCommands: String, CaseIterable, ExpressibleByArgument {
case forward = "forward"
case exit = "exit"
Expand Down
48 changes: 43 additions & 5 deletions Blink/Commands/ssh/SSHPool.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,8 @@ extension SSHPool {

c.localTunnels.forEach { (k, _) in deregister(localForward: k, on: connection) }
c.remoteTunnels.forEach { (k, _) in deregister(remoteForward: k, on: connection) }
c.socks.forEach { (k, _) in deregister(socksBindAddress: k, on: connection) }

// NOTE This is a workaround
c.streams.forEach { (_, s) in s.cancel() }
c.streams = []
Expand Down Expand Up @@ -219,6 +221,34 @@ extension SSHPool {
}
}

// Dynamic Forward
extension SSHPool {
static func register(_ server: SOCKSServer,
bindAddressInfo: OptionalBindAddressInfo,
on connection: SSH.SSHClient) {
let c = control(on: connection)
c?.socks[bindAddressInfo] = server
}

static func deregister(socksBindAddress: OptionalBindAddressInfo, on connection: SSH.SSHClient) {
guard let c = control(on: connection) else {
return
}
if let server = c.socks.removeValue(forKey: socksBindAddress) {
server.close()
}
shared.enforcePersistance(c)
}

static func contains(socksBindAddress: OptionalBindAddressInfo, on connection: SSH.SSHClient) -> Bool {
guard let c = control(on: connection) else {
return false
}

return c.socks[socksBindAddress] != nil
}
}

extension SSHPool {
static func register(stdioStream stream: SSH.Stream, runningCommand command: SSHCommand, on connection: SSH.SSHClient) {
let c = control(on: connection)
Expand Down Expand Up @@ -246,14 +276,15 @@ fileprivate class SSHClientControl {
var numShells: Int = 0
//var shells: [(SSHCommand, SSH.Stream)] = []

var localTunnels: [PortForwardInfo:TunnelControl] = [:]
var remoteTunnels: [PortForwardInfo:TunnelControl] = [:]
var localTunnels: [PortForwardInfo:SSHPortForwardListener] = [:]
var remoteTunnels: [PortForwardInfo:SSHPortForwardClient] = [:]
var socks: [OptionalBindAddressInfo:SOCKSServer] = [:]

var streams: [(SSHCommand, SSH.Stream)] = []

var numChannels: Int {
get {
return numShells + streams.count + localTunnels.count + remoteTunnels.count
return numShells + streams.count + localTunnels.count + remoteTunnels.count + socks.count
}
}

Expand All @@ -276,15 +307,15 @@ fileprivate class SSHClientControl {
return self.host == host && config == self.config ? true : false
}
}

/*
fileprivate protocol TunnelControl {
func close()
}

extension SSHPortForwardListener: TunnelControl {}

extension SSHPortForwardClient: TunnelControl {}

*/
extension PortForwardInfo: Hashable {
public func hash(into hasher: inout Hasher) {
hasher.combine(self.localPort)
Expand All @@ -293,3 +324,10 @@ extension PortForwardInfo: Hashable {
}
}

extension OptionalBindAddressInfo: Hashable {
public func hash(into hasher: inout Hasher) {
hasher.combine(self.bindAddress)
hasher.combine(self.port)
}
}

67 changes: 41 additions & 26 deletions Blink/Commands/ssh/ssh.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
var forwardTunnels: [PortForwardInfo] = []
var remoteTunnels: [PortForwardInfo] = []
var proxyThread: Thread?
var socks: SOCKSServer? = nil
var socks: [OptionalBindAddressInfo] = []

var outStream: DispatchOutputStream?
var inStream: DispatchInputStream?
Expand Down Expand Up @@ -177,7 +177,8 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
}

return self.startInteractiveSessions(conn,
command: cmd,
command: host.remoteCommand,
requestTTY: host.requestTty ?? .auto,
withEnvVars: environment,
sendAgent: host.forwardAgent ?? false)
}
Expand All @@ -189,7 +190,7 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
// TODO Note, we are not merging localForward on host and cmd yet. There can also be -o.
.flatMap { self.startForwardTunnels( (host.localForward ?? []), on: $0, exitOnFailure: host.exitOnForwardFailure ?? false) }
.flatMap { self.startRemoteTunnels( (host.remoteForward ?? []), on: $0, exitOnFailure: host.exitOnForwardFailure ?? false) }
.flatMap { self.startDynamicForwarding($0, command: cmd) }
.flatMap { self.startDynamicForwarding( (host.dynamicForward ?? []), on: $0, exitOnFailure: host.exitOnForwardFailure ?? false) }
.sink(receiveCompletion: { completion in
switch completion {
case .failure(let error):
Expand Down Expand Up @@ -222,10 +223,9 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
if cmd.startsSession { SSHPool.deregister(shellOn: conn) }
forwardTunnels.forEach { SSHPool.deregister(localForward: $0, on: conn) }
remoteTunnels.forEach { SSHPool.deregister(remoteForward: $0, on: conn) }
socks.forEach { SSHPool.deregister(socksBindAddress: $0, on: conn) }
}

self.socks?.close()

return exitCode
}

Expand All @@ -245,31 +245,38 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
}

private func startInteractiveSessions(_ conn: SSH.SSHClient,
command: SSHCommand,
command: String?,
requestTTY: TTYBool,
withEnvVars envVars: [String:String],
sendAgent: Bool) -> SSHConnection {
let rows = Int32(self.device.rows)
let cols = Int32(self.device.cols)
var pty: SSH.SSHClient.PTY? = nil
if command.forceTTY || (self.isTTY && !command.disableTTY && command.command.isEmpty) {
if (requestTTY != .no) &&
(
(requestTTY == .force) ||
// always request a TTY when standard input is a TTY
((requestTTY == .yes) && self.isTTY) ||
// request a TTY when opening a login session
((requestTTY == .auto) && command == nil)
) {
pty = SSH.SSHClient.PTY(rows: rows, columns: cols)
self.device.rawMode = true
}

let session: AnyPublisher<SSH.Stream, Error>

// TERM is explicitely added
var envVars = envVars
envVars["TERM"] = String(cString: getenv("TERM"))

if command.command.isEmpty {
let session: AnyPublisher<SSH.Stream, Error>
if let command = command {
session = conn.requestExec(command: command, withPTY: pty,
withEnvVars: envVars,
withAgentForwarding: sendAgent)
} else {
session = conn.requestInteractiveShell(withPTY: pty,
withEnvVars: envVars,
withAgentForwarding: sendAgent)
} else {
let exec = command.command.joined(separator: " ")
session = conn.requestExec(command: exec, withPTY: pty,
withEnvVars: envVars,
withAgentForwarding: sendAgent)
}

return session.tryMap { s in
Expand Down Expand Up @@ -302,7 +309,7 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
return .just(conn)
}

return conn.requestForward(to: tunnel.bindAddress, port: Int32(tunnel.remotePort),
return conn.requestForward(to: tunnel.bindAddress, port: Int32(tunnel.port),
// Just informative.
from: "stdio", localPort: 22)
.tryMap { s in
Expand Down Expand Up @@ -398,18 +405,26 @@ public func blink_ssh_main(argc: Int32, argv: Argv) -> Int32 {
.eraseToAnyPublisher()
}

private func startDynamicForwarding(_ conn: SSH.SSHClient, command: SSHCommand) -> SSHConnection {
guard let port = command.dynamicForwardingPort else {
private func startDynamicForwarding(_ bindAddresses: [OptionalBindAddressInfo], on conn: SSH.SSHClient, exitOnFailure: Bool) -> SSHConnection {
let bindAddresses = bindAddresses.filter { !SSHPool.contains(socksBindAddress: $0, on: conn) }
if bindAddresses.isEmpty {
return .just(conn)
}

do {
self.socks = try SOCKSServer(port, proxy: conn)
} catch {
return .fail(error: error)
}

return .just(conn)

return bindAddresses.publisher
.flatMap(maxPublishers: .max(1)) { bindAddress -> AnyPublisher<Void, Error> in
do {
let server = try SOCKSServer(bindAddress.port, proxy: conn)
SSHPool.register(server, bindAddressInfo: bindAddress, on: conn)
self.socks.append(bindAddress)
} catch {
return .fail(error: error)
}
return .just(Void())
}
.last()
.map { conn }
.eraseToAnyPublisher()
}

@objc public func sigwinch() {
Expand Down
Loading

0 comments on commit 5a18d59

Please sign in to comment.