Add support for SSL connections to PostgreSQL servers (no unit tests yet)

Author: Andrew Theis

Sources/PostgreSQL/Connection/PostgreSQLConnection.swift

@@ -1,5 +1,6 @@
 import Crypto
 import NIO
+import NIOOpenSSL
 
 /// A PostgreSQL frontend client.
 public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
@@ -127,6 +128,23 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
         /// return the newly enqueued work's future result
         return new
     }
+    
+    /// Trys opening a SLL connection
+    internal func attemptSSLConnection(using tlsConfiguration: TLSConfiguration) -> Future<Void> {
+        return queue.enqueue([.sslSupportRequest(PostgreSQLSSLSupportRequest())]) { message in
+            guard case .sslSupportResponse(let response) = message else {
+                throw PostgreSQLError(identifier: "SSL support check", reason: "Unsupported message encountered during SSL support check: \(message).", source: .capture())
+            }
+            guard response == .supported else {
+                throw PostgreSQLError(identifier: "SSL support check", reason: "tlsConfiguration given in PostgresSQLConfiguration, but SSL connection not supported by PostgreSQL server", source: .capture())
+            }
+            return true
+        }.flatMap {
+            let sslContext = try SSLContext(configuration: tlsConfiguration)
+            let handler = try OpenSSLClientHandler(context: sslContext)
+            return self.channel.pipeline.add(handler: handler, first: true)
+        }
+    }
 
     /// Authenticates the `PostgreSQLClient` using a username with no password.
     public func authenticate(username: String, database: String? = nil, password: String? = nil) -> Future<Void> {

Sources/PostgreSQL/Database/PostgreSQLDatabase.swift

@@ -21,6 +21,11 @@ public final class PostgreSQLDatabase: Database, LogSupporting {
         return Future.flatMap(on: worker) {
             return try PostgreSQLConnection.connect(hostname: config.hostname, port: config.port, on: worker) { error in
                 print("[PostgreSQL] \(error)")
+            }.flatMap(to: PostgreSQLConnection.self) { client in
+                if let tlsConfiguration = config.tlsConfiguration {
+                    return client.attemptSSLConnection(using: tlsConfiguration).transform(to: client)
+                }
+                return Future.map(on: worker) { client }
             }.flatMap(to: PostgreSQLConnection.self) { client in
                 return client.authenticate(
                     username: config.username,

Sources/PostgreSQL/Database/PostgreSQLDatabaseConfig.swift

@@ -1,4 +1,6 @@
 import Foundation
+import NIOOpenSSL
+
 /// Config options for a `PostgreSQLConnection`
 public struct PostgreSQLDatabaseConfig {
     /// Creates a `PostgreSQLDatabaseConfig` with default settings.
@@ -22,38 +24,45 @@ public struct PostgreSQLDatabaseConfig {
     /// Optional password to use for authentication.
     public let password: String?
 
+    /// Optional TLSConfiguration. Set this if your PostgreSQL server requires an SSL configuration
+    public let tlsConfiguration: TLSConfiguration?
+    
     /// Creates a new `PostgreSQLDatabaseConfig`.
-    public init(hostname: String, port: Int = 5432, username: String, database: String? = nil, password: String? = nil) {
+    public init(hostname: String, port: Int = 5432, username: String, database: String? = nil, password: String? = nil, tlsConfiguration: TLSConfiguration? = nil) {
         self.hostname = hostname
         self.port = port
         self.username = username
         self.database = database
         self.password = password
+        self.tlsConfiguration = tlsConfiguration
     }
 
     /// Creates a `PostgreSQLDatabaseConfig` frome a connection string.
-    public init(url: String) throws {
-        guard let urL = URL(string: url),
-            let hostname = urL.host,
-            let port = urL.port,
-            let username = urL.user,
-            let database = URL(string: url)?.path,
-            database.count > 0
-             else {
-                throw PostgreSQLError(identifier: "Bad Connection String",
-                                 reason: "Host could not be parsed",
-                                 possibleCauses: ["Foundation URL is unable to parse the provided connection string"],
-                                 suggestedFixes: ["Check the connection string being passed"],
-                                 source: .capture())
+    public init(url urlString: String, tlsConfiguration: TLSConfiguration? = nil) throws {
+        guard let url = URL(string: urlString),
+            let hostname = url.host,
+            let port = url.port,
+            let username = url.user,
+            url.path.count > 0
+        else {
+            throw PostgreSQLError(
+                identifier: "Bad Connection String",
+                reason: "Host could not be parsed",
+                possibleCauses: ["Foundation URL is unable to parse the provided connection string"],
+                suggestedFixes: ["Check the connection string being passed"],
+                source: .capture()
+            )
         }
         self.hostname = hostname
         self.port = port
         self.username = username
+        let database = url.path
         if database.hasPrefix("/") {
             self.database = database.dropFirst().description
         } else {
             self.database = database
         }
-        self.password = urL.password
+        self.password = url.password
+        self.tlsConfiguration = tlsConfiguration
     }
 }

Sources/PostgreSQL/Message+Parse/PostgreSQLMessageDecoder.swift

@@ -26,8 +26,17 @@ final class PostgreSQLMessageDecoder: ByteToMessageDecoder {
             return .needMoreData
         }
 
-        //// peek at messageSize
+        /// peek at messageSize
         guard let messageSize: Int32 = buffer.peekInteger(skipping: MemoryLayout<Byte>.size) else {
+            // Response from a PostgreSQLSSLSupportRequest will only return a single byte, so we need to handle that case
+            if (messageType == .S || messageType == .N), let data = buffer.readSlice(length: MemoryLayout<Byte>.size) {
+                let decoder = _PostgreSQLMessageDecoder(data: data)
+                let message = try PostgreSQLMessage.sslSupportResponse(decoder.decode())
+                ctx.fireChannelRead(wrapInboundOut(message))
+                VERBOSE("   [message=\(message)]")
+                return .continue
+            }
+            
             VERBOSE("   [needMoreData: messageSize=nil]")
             return .needMoreData
         }

Sources/PostgreSQL/Message+Serialize/PostgreSQLMessageEncoder.swift

@@ -19,6 +19,9 @@ final class PostgreSQLMessageEncoder: MessageToByteEncoder {
         let encoder = _PostgreSQLMessageEncoder()
         let identifier: Byte?
         switch message {
+        case .sslSupportRequest(let request):
+            identifier = nil
+            try request.encode(to: encoder)
         case .startupMessage(let message):
             identifier = nil
             try message.encode(to: encoder)

Sources/PostgreSQL/Message/PostgreSQLErrorReponse.swift renamed to Sources/PostgreSQL/Message/PostgreSQLDiagnosticResponse.swift

@@ -2,6 +2,11 @@ import Bits
 
 /// A frontend or backend PostgreSQL message.
 enum PostgreSQLMessage {
+    /// Asks the server if it supports SSL.
+    case sslSupportRequest(PostgreSQLSSLSupportRequest)
+    /// Response after sending an sslSupportRequest message.
+    case sslSupportResponse(PostgreSQLSSLSupportResponse)
+    /// Startup message
     case startupMessage(PostgreSQLStartupMessage)
     /// Identifies the message as an error.
     case error(PostgreSQLDiagnosticResponse)

Sources/PostgreSQL/Message/PostgreSQLMessage.swift

@@ -0,0 +1,13 @@
+/// A message asking the PostgreSQL server if SSL is supported
+/// For more info, see https://www.postgresql.org/docs/10/static/protocol-flow.html#id-1.10.5.7.11
+struct PostgreSQLSSLSupportRequest: Encodable {
+    /// The SSL request code. The value is chosen to contain 1234 in the most significant 16 bits,
+    /// and 5679 in the least significant 16 bits.
+    let code: Int32 = 80877103
+    
+    /// See Encodable.encode
+    func encode(to encoder: Encoder) throws {
+        var single = encoder.singleValueContainer()
+        try single.encode(code)
+    }
+}

Sources/PostgreSQL/Message/PostgreSQLSSLSupportRequest.swift

@@ -0,0 +1,8 @@
+import Core
+
+/// Response given after sending a PostgreSQLSSLSupportRequest
+/// For more info, see https://www.postgresql.org/docs/10/static/protocol-flow.html#id-1.10.5.7.11
+enum PostgreSQLSSLSupportResponse: UInt8, Decodable {
+    case supported = 0x53
+    case notSupported = 0x4E
+}