Add support for usernames (#72)

fabianfett · web-flow · commit 89a29d457c69 · 2023-07-03T21:23:07.000+02:00
Redis 6.0 adds the ability to specify a username when sending an `AUTH` command. This patch adds this capability to RediStack.
diff --git a/Sources/RediStack/Commands/BasicCommands.swift b/Sources/RediStack/Commands/BasicCommands.swift
@@ -88,6 +88,19 @@ extension RedisClient {
             .map { _ in return () }
     }
 
+    /// Requests the client to authenticate with Redis to allow other commands to be executed.
+    /// - Parameters:
+    ///     - username: The username to authenticate with.
+    ///     - password: The password to authenticate with.
+    ///  Warning: This function should only be used if you are running against Redis 6 or higher.
+    public func authorize(
+        username: String,
+        password: String
+    ) -> EventLoopFuture<Void> {
+        let args = [RESPValue(from: username), RESPValue(from: password)]
+        return self.send(command: "AUTH", with: args).map { _ in return () }
+    }
+
     /// Removes the specified keys. A key is ignored if it does not exist.
     ///
     /// [https://redis.io/commands/del](https://redis.io/commands/del)
diff --git a/Sources/RediStack/Configuration.swift b/Sources/RediStack/Configuration.swift
@@ -80,6 +80,9 @@ extension RedisConnection {
         }
         /// The port of the connection address. If the address is a Unix socket, then it will be `nil`.
         public var port: Int? { self.address.port }
+        /// The user name used to authenticate connections with.
+        /// - Warning: This property should only be provided if you are running against Redis 6 or higher.
+        public let username: String?
         /// The password used to authenticate the connection.
         public let password: String?
         /// The initial database index that the connection should use.
@@ -91,15 +94,17 @@ extension RedisConnection {
 
         /// Creates a new connection configuration with the provided details.
         /// - Parameters:
-        ///     - address: The socket address information to use for creating the Redis connection.
-        ///     - password: The optional password to authenticate the connection with. The default is `nil`.
-        ///     - initialDatabase: The optional database index to initially connect to. The default is `nil`.
-        ///     Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
-        ///     - defaultLogger: The optional prototype logger to use as the default logger instance when generating logs from the connection.
+        ///     - `address`: The socket address information to use for creating the Redis connection.
+        ///     - `username`: The optional username to authenticate the connection with. The default is `nil`.
+        ///     - `password`: The optional password to authenticate the connection with. The default is `nil`.
+        ///     - `initialDatabase`: The optional database index to initially connect to. The default is `nil`.
+        ///       Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
+        ///     - `defaultLogger`: The optional prototype logger to use as the default logger instance when generating logs from the connection.
         ///     If one is not provided, one will be generated. See `RedisLogging.baseConnectionLogger`.
         /// - Throws: `RedisConnection.Configuration.ValidationError` if invalid arguments are provided.
         public init(
             address: SocketAddress,
+            username: String? = nil,
             password: String? = nil,
             initialDatabase: Int? = nil,
             defaultLogger: Logger? = nil
@@ -109,11 +114,36 @@ extension RedisConnection {
             }
 
             self.address = address
+            self.username = username
             self.password = password
             self.initialDatabase = initialDatabase
             self.defaultLogger = defaultLogger ?? Configuration.defaultLogger
         }
 
+        /// Creates a new connection configuration with the provided details.
+        /// - Parameters:
+        ///     - `address`: The socket address information to use for creating the Redis connection.
+        ///     - `password`: The optional password to authenticate the connection with. The default is `nil`.
+        ///     - `initialDatabase`: The optional database index to initially connect to. The default is `nil`.
+        ///       Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
+        ///     - `defaultLogger`: The optional prototype logger to use as the default logger instance when generating logs from the connection.
+        ///     If one is not provided, one will be generated. See `RedisLogging.baseConnectionLogger`.
+        /// - Throws: `RedisConnection.Configuration.ValidationError` if invalid arguments are provided.
+        public init(
+            address: SocketAddress,
+            password: String? = nil,
+            initialDatabase: Int? = nil,
+            defaultLogger: Logger? = nil
+        ) throws {
+            try self.init(
+                address: address,
+                username: nil,
+                password: password,
+                initialDatabase: initialDatabase,
+                defaultLogger: defaultLogger
+            )
+        }
+
         /// Creates a new connection configuration with exact details.
         /// - Parameters:
         ///     - hostname: The remote hostname to connect to.
@@ -214,6 +244,9 @@ extension RedisConnectionPool {
         // this needs to be var so it can be updated by the pool with the pool id
         /// The logger prototype that will be used by connections by default when generating logs.
         public internal(set) var connectionDefaultLogger: Logger
+        /// The username used to authenticate connections.
+        /// - Warning: This property should only be provided if you are running against Redis 6 or higher.
+        public let connectionUsername: String?
         /// The password used to authenticate connections.
         public let connectionPassword: String?
         /// The initial database index that connections should use.
@@ -224,7 +257,7 @@ extension RedisConnectionPool {
         /// Creates a new connection factory configuration with the provided options.
         /// - Parameters:
         ///     - connectionInitialDatabase: The optional database index to initially connect to. The default is `nil`.
-        ///     Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
+        ///       Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
         ///     - connectionPassword: The optional password to authenticate connections with. The default is `nil`.
         ///     - connectionDefaultLogger: The optional prototype logger to use as the default logger instance when generating logs from connections.
         ///     If one is not provided, one will be generated. See `RedisLogging.baseConnectionLogger`.
@@ -234,8 +267,34 @@ extension RedisConnectionPool {
             connectionPassword: String? = nil,
             connectionDefaultLogger: Logger? = nil,
             tcpClient: ClientBootstrap? = nil
+        ) {
+            self.init(
+                connectionInitialDatabase: connectionInitialDatabase,
+                connectionUsername: nil,
+                connectionPassword: connectionPassword,
+                connectionDefaultLogger: connectionDefaultLogger,
+                tcpClient: tcpClient
+            )
+        }
+
+        /// Creates a new connection factory configuration with the provided options.
+        /// - Parameters:
+        ///     - connectionInitialDatabase: The optional database index to initially connect to. The default is `nil`.
+        ///       Redis by default opens connections against index `0`, so only set this value if the desired default is not `0`.
+        ///     - connectionUsername: The optional username to authenticate connections with. The default is `nil`. Works only with Redis 6 and greater.
+        ///     - connectionPassword: The optional password to authenticate connections with. The default is `nil`.
+        ///     - connectionDefaultLogger: The optional prototype logger to use as the default logger instance when generating logs from connections.
+        ///       If one is not provided, one will be generated. See `RedisLogging.baseConnectionLogger`.
+        ///     - tcpClient: If you have chosen to configure a `NIO.ClientBootstrap` yourself, this will be used instead of the `.makeRedisTCPClient` factory instance.
+        public init(
+            connectionInitialDatabase: Int? = nil,
+            connectionUsername: String? = nil,
+            connectionPassword: String? = nil,
+            connectionDefaultLogger: Logger? = nil,
+            tcpClient: ClientBootstrap? = nil
         ) {
             self.connectionInitialDatabase = connectionInitialDatabase
+            self.connectionUsername = connectionUsername
             self.connectionPassword = connectionPassword
             self.connectionDefaultLogger = connectionDefaultLogger ?? RedisConnection.Configuration.defaultLogger
             self.tcpClient = tcpClient
diff --git a/Sources/RediStack/RedisConnection.swift b/Sources/RediStack/RedisConnection.swift
@@ -57,25 +57,12 @@ extension RedisConnection {
     ) -> EventLoopFuture<RedisConnection> {
         let client = client ?? .makeRedisTCPClient(group: eventLoop)
 
-        var future = client
+        return client
             .connect(to: config.address)
-            .map { return RedisConnection(configuredRESPChannel: $0, backgroundLogger: config.defaultLogger) }
-
-        // if a password is specified, use it to authenticate before further operations happen
-        if let password = config.password {
-            future = future.flatMap { connection in
-                return connection.authorize(with: password).map { connection }
-            }
-        }
-
-        // if a database index is specified, use it to switch the selected database before further operations happen
-        if let database = config.initialDatabase {
-            future = future.flatMap { connection in
-                return connection.select(database: database).map { connection }
+            .flatMap {
+                let connection = RedisConnection(configuredRESPChannel: $0, backgroundLogger: config.defaultLogger)
+                return connection.start(configuration: config).map({ _ in connection })
             }
-        }
-
-        return future
     }
 }
 
@@ -194,6 +181,28 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
         self.logger.trace("connection created")
     }
 
+    func start(configuration: Configuration) -> EventLoopFuture<Void> {
+        let future: EventLoopFuture<Void>
+
+        // if a password is specified, use it to authenticate before further operations happen
+        if let password = configuration.password {
+            if let username = configuration.username {
+                future = self.authorize(username: username, password: password)
+            } else {
+                future = self.authorize(with: password)
+            }
+        } else {
+            future = self.eventLoop.makeSucceededVoidFuture()
+        }
+
+        // if a database index is specified, use it to switch the selected database before further operations happen
+        if let database = configuration.initialDatabase {
+            return future.flatMap { self.select(database: database) }
+        }
+
+        return future
+    }
+
     internal enum ConnectionState {
         case open
         case pubsub(RedisPubSubHandler)
@@ -223,15 +232,27 @@ extension RedisConnection {
     /// - Returns: A `NIO.EventLoopFuture` that resolves with the command's result stored in a `RESPValue`.
     ///     If a `RedisError` is returned, the future will be failed instead.
     public func send(command: String, with arguments: [RESPValue]) -> EventLoopFuture<RESPValue> {
-        self.eventLoop.flatSubmit {
-            return self.send(command: command, with: arguments, logger: nil)
-        }
+        return self.send(command: command, with: arguments, logger: nil)
     }
 
     internal func send(
         command: String,
         with arguments: [RESPValue],
         logger: Logger?
+    ) -> EventLoopFuture<RESPValue> {
+        if self.eventLoop.inEventLoop {
+            return self.send0(command: command, with: arguments, logger: logger)
+        }
+
+        return self.eventLoop.flatSubmit {
+            self.send0(command: command, with: arguments, logger: logger)
+        }
+    }
+
+    private func send0(
+        command: String,
+        with arguments: [RESPValue],
+        logger: Logger?
     ) -> EventLoopFuture<RESPValue> {
         self.eventLoop.preconditionInEventLoop()
 
diff --git a/Sources/RediStack/RedisConnectionPool.swift b/Sources/RediStack/RedisConnectionPool.swift
@@ -220,6 +220,7 @@ extension RedisConnectionPool {
         do {
             connectionConfig = try .init(
                 address: nextTarget,
+                username: factoryConfig.connectionUsername,
                 password: factoryConfig.connectionPassword,
                 initialDatabase: factoryConfig.connectionInitialDatabase,
                 defaultLogger: factoryConfig.connectionDefaultLogger
diff --git a/Tests/RediStackTests/RedisConnectionTests.swift b/Tests/RediStackTests/RedisConnectionTests.swift
@@ -20,16 +20,16 @@ import XCTest
 
 final class RedisConnectionTests: XCTestCase {
 
-}
+    var logger: Logger {
+        Logger(label: "RedisConnectionTests")
+    }
 
-// MARK: Unexpected Closures
-extension RedisConnectionTests {
     func test_connectionUnexpectedlyCloses_invokesCallback() throws {
         let loop = EmbeddedEventLoop()
 
         let expectedClosureConnection = RedisConnection(
             configuredRESPChannel: EmbeddedChannel(loop: loop),
-            backgroundLogger: Logger(label: "")
+            backgroundLogger: self.logger
         )
         let expectedClosureExpectation = self.expectation(description: "this should not be fulfilled")
         expectedClosureExpectation.isInverted = true
@@ -49,4 +49,56 @@ extension RedisConnectionTests {
 
         self.waitForExpectations(timeout: 0.5)
     }
+
+    func testAuthorizationWithUsername() {
+        var maybeSocketAddress: SocketAddress?
+        XCTAssertNoThrow(maybeSocketAddress = try SocketAddress.makeAddressResolvingHost("localhost", port: 0))
+        guard let socketAddress = maybeSocketAddress else { return XCTFail("Expected a socketAddress") }
+        var maybeConfiguration: RedisConnection.Configuration?
+        XCTAssertNoThrow(maybeConfiguration = try .init(address: socketAddress, username: "username", password: "password"))
+        guard let configuration = maybeConfiguration else { return XCTFail("Expected a configuration") }
+
+        let channel = EmbeddedChannel(handlers: [RedisCommandHandler()])
+        XCTAssertNoThrow(try channel.connect(to: socketAddress).wait())
+
+        let connection = RedisConnection(configuredRESPChannel: channel, backgroundLogger: self.logger)
+        let future = connection.start(configuration: configuration)
+
+        var outgoing: RESPValue?
+        XCTAssertNoThrow(outgoing = try channel.readOutbound(as: RESPValue.self))
+        XCTAssertEqual(outgoing, .array([.bulkString("AUTH"), .bulkString("username"), .bulkString("password")]))
+        XCTAssertNoThrow(try channel.writeInbound(RESPValue.simpleString("OK")))
+        XCTAssertNoThrow(try future.wait())
+    }
+
+    func testAuthorizationWithoutUsername() {
+        var maybeSocketAddress: SocketAddress?
+        XCTAssertNoThrow(maybeSocketAddress = try SocketAddress.makeAddressResolvingHost("localhost", port: 0))
+        guard let socketAddress = maybeSocketAddress else { return XCTFail("Expected a socketAddress") }
+        var maybeConfiguration: RedisConnection.Configuration?
+        XCTAssertNoThrow(maybeConfiguration = try .init(address: socketAddress, password: "password"))
+        guard let configuration = maybeConfiguration else { return XCTFail("Expected a configuration") }
+
+        let channel = EmbeddedChannel(handlers: [RedisCommandHandler()])
+        XCTAssertNoThrow(try channel.connect(to: socketAddress).wait())
+
+        let connection = RedisConnection(configuredRESPChannel: channel, backgroundLogger: self.logger)
+        let future = connection.start(configuration: configuration)
+
+        var outgoing: RESPValue?
+        XCTAssertNoThrow(outgoing = try channel.readOutbound(as: RESPValue.self))
+        XCTAssertEqual(outgoing, .array([.bulkString("AUTH"), .bulkString("password")]))
+        XCTAssertNoThrow(try channel.writeInbound(RESPValue.simpleString("OK")))
+        XCTAssertNoThrow(try future.wait())
+    }
+}
+
+extension RESPValue {
+    static func bulkString(_ string: String) -> Self {
+        .bulkString(ByteBuffer(string: string))
+    }
+
+    static func simpleString(_ string: String) -> Self {
+        .simpleString(ByteBuffer(string: string))
+    }
 }
