Updated Jim Studt's PR

apple · Nov 19, 2022 · f5e22a8 · f5e22a8
1 parent c905128
commit f5e22a8
Show file tree

Hide file tree

Showing 19 changed files with 83 additions and 6 deletions.
diff --git a/Sources/NIOSSH/Child Channels/ChildChannelOptions.swift b/Sources/NIOSSH/Child Channels/ChildChannelOptions.swift
@@ -29,6 +29,9 @@ public struct SSHChildChannelOptions {
 
     /// See: ``SSHChildChannelOptions/Types/PeerMaximumMessageLengthOption``.
     public static let peerMaximumMessageLength: SSHChildChannelOptions.Types.PeerMaximumMessageLengthOption = .init()
+
+    /// - seealso: `UsernameOption`.
+    public static let username: SSHChildChannelOptions.Types.UsernameOption = .init()
 }
 
 extension SSHChildChannelOptions {
@@ -61,7 +64,14 @@ extension SSHChildChannelOptions.Types {
     /// ``SSHChildChannelOptions/Types/PeerMaximumMessageLengthOption`` allows users to query the maximum packet size value reported by the remote peer for a given channel.
     public struct PeerMaximumMessageLengthOption: ChannelOption, Sendable {
         public typealias Value = UInt32
-
+
+        public init() {}
+    }
+
+    /// `UsernameOption` allows users to query the authenticated username of the channel.
+    public struct UsernameOption: ChannelOption {
+        public typealias Value = String?
+
         public init() {}
     }
 }
diff --git a/Sources/NIOSSH/Child Channels/SSHChannelMultiplexer.swift b/Sources/NIOSSH/Child Channels/SSHChannelMultiplexer.swift
@@ -93,6 +93,9 @@ extension SSHChannelMultiplexer {
             self.erroredChannels.append(channelID)
         }
     }
+
+    // The username which the server accepted in authorization
+    var username: String? { delegate?.username }
 }
 
 // MARK: Calls from SSH handlers.
@@ -218,6 +221,8 @@ extension SSHChannelMultiplexer {
 protocol SSHMultiplexerDelegate {
     var channel: Channel? { get }
 
+    var username: String? { get }
+
     func writeFromChildChannel(_: SSHMessage, _: EventLoopPromise<Void>?)
 
     func flushFromChildChannel()

diff --git a/Sources/NIOSSH/Child Channels/SSHChildChannel.swift b/Sources/NIOSSH/Child Channels/SSHChildChannel.swift
@@ -228,6 +228,8 @@ extension SSHChildChannel: Channel, ChannelCore {
             return self.type! as! Option.Value
         case _ as SSHChildChannelOptions.Types.PeerMaximumMessageLengthOption:
             return self.peerMaxMessageSize as! Option.Value
+        case _ as SSHChildChannelOptions.Types.UsernameOption:
+            return multiplexer.username as! Option.Value
         case _ as ChannelOptions.Types.AutoReadOption:
             return self.autoRead as! Option.Value
         case _ as ChannelOptions.Types.AllowRemoteHalfClosureOption:

diff --git a/Sources/NIOSSH/Connection State Machine/Operations/AcceptsUserAuthMessages.swift b/Sources/NIOSSH/Connection State Machine/Operations/AcceptsUserAuthMessages.swift
@@ -15,6 +15,8 @@
 protocol AcceptsUserAuthMessages {
     var userAuthStateMachine: UserAuthenticationStateMachine { get set }
 
+    var connectionAttributes: SSHConnectionStateMachine.Attributes? { get }
+
     var role: SSHConnectionRole { get }
 }
 
@@ -60,14 +62,15 @@ extension AcceptsUserAuthMessages {
 
     mutating func receiveUserAuthRequest(_ message: SSHMessage.UserAuthRequestMessage) throws -> SSHConnectionStateMachine.StateMachineInboundProcessResult {
         let result = try self.userAuthStateMachine.receiveUserAuthRequest(message)
-
+        
         if let future = result {
             var banner: SSHServerConfiguration.UserAuthBanner?
             if case .server(let config) = role {
                 banner = config.banner
             }
 
-            return .possibleFutureMessage(future.map { Self.transform($0, banner: banner) })
+            let connectionAttributes = self.connectionAttributes
+            return .possibleFutureMessage(future.map { Self.transform($0, connectionAttributes: connectionAttributes, username: message.username, banner: banner) })
         } else {
             return .noMessage
         }
@@ -96,9 +99,11 @@ extension AcceptsUserAuthMessages {
         return .event(NIOUserAuthBannerEvent(message: message.message, languageTag: message.languageTag))
     }
 
-    private static func transform(_ result: NIOSSHUserAuthenticationResponseMessage, banner: SSHServerConfiguration.UserAuthBanner? = nil) -> SSHMultiMessage {
+    private static func transform(_ result: NIOSSHUserAuthenticationResponseMessage, connectionAttributes: SSHConnectionStateMachine.Attributes?, username: String, banner: SSHServerConfiguration.UserAuthBanner? = nil) -> SSHMultiMessage {
         switch result {
         case .success:
+            connectionAttributes?.username = username
+
             if let banner = banner {
                 // Send banner bundled with auth success to avoid leaking any information to unauthenticated clients.
                 // Note that this is by no means the only option according to RFC 4252

diff --git a/Sources/NIOSSH/Connection State Machine/SSHConnectionStateMachine.swift b/Sources/NIOSSH/Connection State Machine/SSHConnectionStateMachine.swift
@@ -57,11 +57,21 @@ struct SSHConnectionStateMachine {
         case sentDisconnect(SSHConnectionRole)
     }
 
+    class Attributes {
+        var username: String? = nil
+    }
+
     /// The state of this state machine.
     private var state: State
+
+    /// Attributes of the connection which can be changed by messages handlers
+    private let attributes: Attributes
+
+    var username: String? { attributes.username }
 
     init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type] = Constants.bundledTransportProtectionSchemes) {
-        self.state = .idle(IdleState(role: role, protectionSchemes: protectionSchemes))
+        self.attributes = .init()
+        self.state = .idle(IdleState(role: role, protectionSchemes: protectionSchemes, attributes: attributes))
     }
 
     func start() -> SSHMultiMessage? {

diff --git a/Sources/NIOSSH/Connection State Machine/States/ActiveState.swift b/Sources/NIOSSH/Connection State Machine/States/ActiveState.swift
@@ -31,13 +31,16 @@ extension SSHConnectionStateMachine {
 
         internal var sessionIdentifier: ByteBuffer
 
+        internal weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previous: UserAuthenticationState) {
             self.role = previous.role
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
             self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
+            self.connectionAttributes = previous.connectionAttributes
         }
 
         init(_ previous: RekeyingReceivedNewKeysState) {
@@ -47,6 +50,7 @@ extension SSHConnectionStateMachine {
             self.remoteVersion = previous.remoteVersion
             self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
+            self.connectionAttributes = previous.connectionAttributes
         }
 
         init(_ previous: RekeyingSentNewKeysState) {
@@ -56,6 +60,7 @@ extension SSHConnectionStateMachine {
             self.remoteVersion = previous.remoteVersion
             self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
+            self.connectionAttributes = previous.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/IdleState.swift b/Sources/NIOSSH/Connection State Machine/States/IdleState.swift
@@ -23,10 +23,13 @@ extension SSHConnectionStateMachine {
 
         internal var protectionSchemes: [NIOSSHTransportProtection.Type]
 
-        init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type]) {
+        internal weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
+        init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type], attributes: SSHConnectionStateMachine.Attributes) {
             self.role = role
             self.serializer = SSHPacketSerializer()
             self.protectionSchemes = protectionSchemes
+            self.connectionAttributes = attributes
         }
     }
 }
diff --git a/Sources/NIOSSH/Connection State Machine/States/KeyExchangeState.swift b/Sources/NIOSSH/Connection State Machine/States/KeyExchangeState.swift
@@ -33,13 +33,16 @@ extension SSHConnectionStateMachine {
         /// The backing state machine.
         var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(sentVersionState state: SentVersionState, allocator: ByteBufferAllocator, loop: EventLoop, remoteVersion: String) {
             self.role = state.role
             self.parser = state.parser
             self.serializer = state.serializer
             self.remoteVersion = remoteVersion
             self.protectionSchemes = state.protectionSchemes
             self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: state.role, remoteVersion: remoteVersion, protectionSchemes: state.protectionSchemes, previousSessionIdentifier: nil)
+            self.connectionAttributes = state.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/ReceivedKexInitWhenActiveState.swift b/Sources/NIOSSH/Connection State Machine/States/ReceivedKexInitWhenActiveState.swift
@@ -34,6 +34,8 @@ extension SSHConnectionStateMachine {
 
         internal var sessionIdentifier: ByteBuffer
 
+        internal weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previous: ActiveState, allocator: ByteBufferAllocator, loop: EventLoop) {
             self.role = previous.role
             self.serializer = previous.serializer
@@ -42,6 +44,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
             self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: previous.role, remoteVersion: previous.remoteVersion, protectionSchemes: previous.protectionSchemes, previousSessionIdentifier: self.sessionIdentifier)
+            self.connectionAttributes = previous.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/ReceivedNewKeysState.swift b/Sources/NIOSSH/Connection State Machine/States/ReceivedNewKeysState.swift
@@ -39,6 +39,8 @@ extension SSHConnectionStateMachine {
         /// The user auth state machine that drives user authentication.
         var userAuthStateMachine: UserAuthenticationStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(keyExchangeState state: KeyExchangeState,
              loop: EventLoop) {
             self.role = state.role
@@ -53,6 +55,7 @@ extension SSHConnectionStateMachine {
             self.userAuthStateMachine = UserAuthenticationStateMachine(role: self.role,
                                                                        loop: loop,
                                                                        sessionID: self.sessionIdentifier)
+            self.connectionAttributes = state.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/RekeyingReceivedNewKeysState.swift b/Sources/NIOSSH/Connection State Machine/States/RekeyingReceivedNewKeysState.swift
@@ -36,6 +36,8 @@ extension SSHConnectionStateMachine {
         /// The backing state machine.
         var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previousState: RekeyingState) {
             self.role = previousState.role
             self.parser = previousState.parser
@@ -44,6 +46,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
+            self.connectionAttributes = previousState.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/RekeyingSentNewKeysState.swift b/Sources/NIOSSH/Connection State Machine/States/RekeyingSentNewKeysState.swift
@@ -36,6 +36,8 @@ extension SSHConnectionStateMachine {
         /// The backing state machine.
         var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previousState: RekeyingState) {
             self.role = previousState.role
             self.parser = previousState.parser
@@ -44,6 +46,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
+            self.connectionAttributes = previousState.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/RekeyingState.swift b/Sources/NIOSSH/Connection State Machine/States/RekeyingState.swift
@@ -35,6 +35,8 @@ extension SSHConnectionStateMachine {
         /// The backing state machine.
         var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previousState: ReceivedKexInitWhenActiveState) {
             self.role = previousState.role
             self.parser = previousState.parser
@@ -43,6 +45,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
+            self.connectionAttributes = previousState.connectionAttributes
         }
 
         init(_ previousState: SentKexInitWhenActiveState) {
@@ -53,6 +56,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentitifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
+            self.connectionAttributes = previousState.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/SentKexInitWhenActiveState.swift b/Sources/NIOSSH/Connection State Machine/States/SentKexInitWhenActiveState.swift
@@ -34,6 +34,8 @@ extension SSHConnectionStateMachine {
 
         internal var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
+        internal weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(_ previous: ActiveState, allocator: ByteBufferAllocator, loop: EventLoop) {
             self.role = previous.role
             self.serializer = previous.serializer
@@ -42,6 +44,7 @@ extension SSHConnectionStateMachine {
             self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentitifier = previous.sessionIdentifier
             self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: self.role, remoteVersion: self.remoteVersion, protectionSchemes: self.protectionSchemes, previousSessionIdentifier: previous.sessionIdentifier)
+            self.connectionAttributes = previous.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/SentNewKeysState.swift b/Sources/NIOSSH/Connection State Machine/States/SentNewKeysState.swift
@@ -39,6 +39,8 @@ extension SSHConnectionStateMachine {
         /// The user auth state machine that drives user authentication.
         var userAuthStateMachine: UserAuthenticationStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(keyExchangeState state: KeyExchangeState,
              loop: EventLoop) {
             self.role = state.role
@@ -53,6 +55,7 @@ extension SSHConnectionStateMachine {
             self.userAuthStateMachine = UserAuthenticationStateMachine(role: self.role,
                                                                        loop: loop,
                                                                        sessionID: self.sessionIdentifier)
+            self.connectionAttributes = state.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/SentVersionState.swift b/Sources/NIOSSH/Connection State Machine/States/SentVersionState.swift
@@ -28,6 +28,8 @@ extension SSHConnectionStateMachine {
 
         var protectionSchemes: [NIOSSHTransportProtection.Type]
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         private let allocator: ByteBufferAllocator
 
         init(idleState state: IdleState, allocator: ByteBufferAllocator) {
@@ -37,6 +39,7 @@ extension SSHConnectionStateMachine {
 
             self.parser = SSHPacketParser(allocator: allocator)
             self.allocator = allocator
+            self.connectionAttributes = state.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/Connection State Machine/States/UserAuthenticationState.swift b/Sources/NIOSSH/Connection State Machine/States/UserAuthenticationState.swift
@@ -35,6 +35,8 @@ extension SSHConnectionStateMachine {
         /// The backing state machine.
         var userAuthStateMachine: UserAuthenticationStateMachine
 
+        weak var connectionAttributes: SSHConnectionStateMachine.Attributes?
+
         init(sentNewKeysState state: SentNewKeysState) {
             self.role = state.role
             self.parser = state.parser
@@ -43,6 +45,7 @@ extension SSHConnectionStateMachine {
             self.remoteVersion = state.remoteVersion
             self.protectionSchemes = state.protectionSchemes
             self.sessionIdentifier = state.sessionIdentifier
+            self.connectionAttributes = state.connectionAttributes
         }
 
         init(receivedNewKeysState state: ReceivedNewKeysState) {
@@ -53,6 +56,7 @@ extension SSHConnectionStateMachine {
             self.remoteVersion = state.remoteVersion
             self.protectionSchemes = state.protectionSchemes
             self.sessionIdentifier = state.sessionIdentifier
+            self.connectionAttributes = state.connectionAttributes
         }
     }
 }

diff --git a/Sources/NIOSSH/NIOSSHHandler.swift b/Sources/NIOSSH/NIOSSHHandler.swift
@@ -57,6 +57,9 @@ public final class NIOSSHHandler {
 
     private var pendingGlobalRequestResponses: CircularBuffer<PendingGlobalRequestResponse?>
 
+    // The authenticated username, if there was one.
+    var username: String? { stateMachine.username }
+
     /// Construct a new ``NIOSSHHandler``.
     ///
     /// - parameters:

diff --git a/Tests/NIOSSHTests/ChildChannelMultiplexerTests.swift b/Tests/NIOSSHTests/ChildChannelMultiplexerTests.swift
@@ -22,6 +22,8 @@ import XCTest
 /// This reduces the testing surface area somewhat, which greatly helps us to test the
 /// implementation of the multiplexer and child channels.
 final class DummyDelegate: SSHMultiplexerDelegate {
+    var username : String? = "dummy"
+
     var _channel: EmbeddedChannel = EmbeddedChannel()
 
     var writes: MarkedCircularBuffer<(SSHMessage, EventLoopPromise<Void>?)> = MarkedCircularBuffer(initialCapacity: 8)