Skip to content

State machine #135

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 33 commits into from
Feb 24, 2021
Merged
Show file tree
Hide file tree
Changes from 29 commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
6cf80c0
Adds PSQLFrontendMessage & PSQLBackendMessage
fabianfett Feb 2, 2021
da88bf6
State machine
fabianfett Feb 3, 2021
988c746
Removed Xcode headers.
fabianfett Feb 7, 2021
d451940
Apply suggestions from code review
fabianfett Feb 7, 2021
47e0d4d
Code review
fabianfett Feb 7, 2021
3da358e
Apply suggestions from code review
fabianfett Feb 7, 2021
bf898d9
Code review
fabianfett Feb 7, 2021
8c7082a
Apply suggestions from code review
fabianfett Feb 7, 2021
333c2ea
Code review
fabianfett Feb 9, 2021
02a143d
Add rudementary sasl support
fabianfett Feb 11, 2021
589b6fb
Merge branch 'master' into ff-state-machine
fabianfett Feb 12, 2021
ab1aed3
Update Sources/PostgresNIO/Connection/PostgresConnection+Notification…
fabianfett Feb 12, 2021
c6cca43
Code review
fabianfett Feb 12, 2021
63abdf3
Code review
fabianfett Feb 12, 2021
a416292
A little more error handling
fabianfett Feb 12, 2021
f857a43
Error handling
fabianfett Feb 16, 2021
39ccea0
Better logging
fabianfett Feb 16, 2021
8bb4141
Fixes!
fabianfett Feb 16, 2021
00b4c5e
Some better state handling when closing
fabianfett Feb 17, 2021
005df8c
State machine tests
fabianfett Feb 18, 2021
5d64ea0
Better cleanup in error states
fabianfett Feb 18, 2021
e510315
Merge branch 'master' into ff-state-machine
fabianfett Feb 18, 2021
56ae38b
Cherry pick to be reverted.
fabianfett Feb 18, 2021
099fbaa
PreparedStatementStateMachine tests
fabianfett Feb 18, 2021
eac4465
Code review
fabianfett Feb 18, 2021
a001aea
Enable trace logging to better find the flaky tests
fabianfett Feb 19, 2021
559a12e
PSQLChannelHandler logging + cleanup
fabianfett Feb 19, 2021
9cafe17
Merge branch 'master' into ff-state-machine
fabianfett Feb 20, 2021
f2c7a61
PR review
fabianfett Feb 22, 2021
580628d
Code review
fabianfett Feb 23, 2021
e7b9f72
Update Sources/PostgresNIO/New/Connection State Machine/Authenticatio…
fabianfett Feb 23, 2021
a6fd040
Last code comment
fabianfett Feb 23, 2021
94399f1
Last code comment fix
fabianfett Feb 24, 2021
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -99,9 +99,10 @@ jobs:
run: swift test --enable-test-discovery --sanitize=thread
env:
POSTGRES_HOSTNAME: psql
POSTGRES_USERNAME: vapor_username
POSTGRES_USER: vapor_username
POSTGRES_DB: vapor_database
POSTGRES_PASSWORD: vapor_password
POSTGRES_DATABASE: vapor_database
POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }}

# Run package tests on macOS against supported PSQL versions
macos:
Expand Down Expand Up @@ -138,6 +139,7 @@ jobs:
run: swift test --enable-test-discovery --sanitize=thread
env:
POSTGRES_HOSTNAME: 127.0.0.1
POSTGRES_USERNAME: vapor_username
POSTGRES_USER: vapor_username
POSTGRES_DB: postgres
POSTGRES_PASSWORD: vapor_password
POSTGRES_DATABASE: postgres
POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }}
2 changes: 2 additions & 0 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ let package = Package(
.product(name: "Logging", package: "swift-log"),
.product(name: "Metrics", package: "swift-metrics"),
.product(name: "NIO", package: "swift-nio"),
.product(name: "NIOTLS", package: "swift-nio"),
.product(name: "NIOFoundationCompat", package: "swift-nio"),
.product(name: "NIOSSL", package: "swift-nio-ssl"),
]),
.testTarget(name: "PostgresNIOTests", dependencies: [
Expand Down
158 changes: 9 additions & 149 deletions Sources/PostgresNIO/Connection/PostgresConnection+Authenticate.swift
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import Crypto
import NIO
import Logging

extension PostgresConnection {
public func authenticate(
Expand All @@ -9,155 +7,17 @@ extension PostgresConnection {
password: String? = nil,
logger: Logger = .init(label: "codes.vapor.postgres")
) -> EventLoopFuture<Void> {
let auth = PostgresAuthenticationRequest(
let authContext = AuthContext(
username: username,
database: database,
password: password
)
return self.send(auth, logger: self.logger)
}
}

// MARK: Private
password: password,
database: database)
let outgoing = PSQLOutgoingEvent.authenticate(authContext)
self.underlying.channel.triggerUserOutboundEvent(outgoing, promise: nil)

private final class PostgresAuthenticationRequest: PostgresRequest {
enum State {
case ready
case saslInitialSent(SASLAuthenticationManager<SASLMechanism.SCRAM.SHA256>)
case saslChallengeResponse(SASLAuthenticationManager<SASLMechanism.SCRAM.SHA256>)
case saslWaitOkay
case done
}

let username: String
let database: String?
let password: String?
var state: State

init(username: String, database: String?, password: String?) {
self.state = .ready
self.username = username
self.database = database
self.password = password
}

func log(to logger: Logger) {
logger.debug("Logging into Postgres db \(self.database ?? "nil") as \(self.username)")
}

func respond(to message: PostgresMessage) throws -> [PostgresMessage]? {
if case .error = message.identifier {
// terminate immediately on error
return nil
}

switch self.state {
case .ready:
switch message.identifier {
case .authentication:
let auth = try PostgresMessage.Authentication(message: message)
switch auth {
case .md5(let salt):
let pwdhash = self.md5((self.password ?? "") + self.username).hexdigest()
let hash = "md5" + self.md5(self.bytes(pwdhash) + salt).hexdigest()
return try [PostgresMessage.Password(string: hash).message()]
case .plaintext:
return try [PostgresMessage.Password(string: self.password ?? "").message()]
case .saslMechanisms(let saslMechanisms):
if saslMechanisms.contains("SCRAM-SHA-256") && self.password != nil {
let saslManager = SASLAuthenticationManager(asClientSpeaking:
SASLMechanism.SCRAM.SHA256(username: self.username, password: { self.password! }))
var message: PostgresMessage?

if (try saslManager.handle(message: nil, sender: { bytes in
message = try PostgresMessage.SASLInitialResponse(mechanism: "SCRAM-SHA-256", initialData: bytes).message()
})) {
self.state = .saslWaitOkay
} else {
self.state = .saslInitialSent(saslManager)
}
return [message].compactMap { $0 }
} else {
throw PostgresError.protocol("Unable to authenticate with any available SASL mechanism: \(saslMechanisms)")
}
case .saslContinue, .saslFinal:
throw PostgresError.protocol("Unexpected SASL response to start message: \(message)")
case .ok:
self.state = .done
return []
}
default: throw PostgresError.protocol("Unexpected response to start message: \(message)")
}
case .saslInitialSent(let manager),
.saslChallengeResponse(let manager):
switch message.identifier {
case .authentication:
let auth = try PostgresMessage.Authentication(message: message)
switch auth {
case .saslContinue(let data), .saslFinal(let data):
var message: PostgresMessage?
if try manager.handle(message: data, sender: { bytes in
message = try PostgresMessage.SASLResponse(responseData: bytes).message()
}) {
self.state = .saslWaitOkay
} else {
self.state = .saslChallengeResponse(manager)
}
return [message].compactMap { $0 }
default: throw PostgresError.protocol("Unexpected response during SASL negotiation: \(message)")
}
default: throw PostgresError.protocol("Unexpected response during SASL negotiation: \(message)")
}
case .saslWaitOkay:
switch message.identifier {
case .authentication:
let auth = try PostgresMessage.Authentication(message: message)
switch auth {
case .ok:
self.state = .done
return []
default: throw PostgresError.protocol("Unexpected response while waiting for post-SASL ok: \(message)")
}
default: throw PostgresError.protocol("Unexpected response while waiting for post-SASL ok: \(message)")
}
case .done:
switch message.identifier {
case .parameterStatus:
// self.status[status.parameter] = status.value
return []
case .backendKeyData:
// self.processID = data.processID
// self.secretKey = data.secretKey
return []
case .readyForQuery:
return nil
default: throw PostgresError.protocol("Unexpected response to password authentication: \(message)")
}
return self.underlying.channel.pipeline.handler(type: PSQLEventsHandler.self).flatMap { handler in
handler.authenticateFuture
}.flatMapErrorThrowing { error in
throw error.asAppropriatePostgresError
}

}

func start() throws -> [PostgresMessage] {
return try [
PostgresMessage.Startup.versionThree(parameters: [
"user": self.username,
"database": self.database ?? username
]).message()
]
}

// MARK: Private

private func md5(_ string: String) -> [UInt8] {
return md5(self.bytes(string))
}

private func md5(_ message: [UInt8]) -> [UInt8] {
let digest = Insecure.MD5.hash(data: message)
return .init(digest)
}

func bytes(_ string: String) -> [UInt8] {
return Array(string.utf8)
}
}
62 changes: 20 additions & 42 deletions Sources/PostgresNIO/Connection/PostgresConnection+Connect.swift
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import Logging
import NIO

extension PostgresConnection {
Expand All @@ -9,47 +8,26 @@ extension PostgresConnection {
logger: Logger = .init(label: "codes.vapor.postgres"),
on eventLoop: EventLoop
) -> EventLoopFuture<PostgresConnection> {
let bootstrap = ClientBootstrap(group: eventLoop)
.channelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
return bootstrap.connect(to: socketAddress).flatMap { channel in
return channel.pipeline.addHandlers([
ByteToMessageHandler(PostgresMessageDecoder(logger: logger)),
MessageToByteHandler(PostgresMessageEncoder(logger: logger)),
PostgresRequestHandler(logger: logger),
PostgresErrorHandler(logger: logger)
]).map {
return PostgresConnection(channel: channel, logger: logger)
}
}.flatMap { (conn: PostgresConnection) in
if let tlsConfiguration = tlsConfiguration {
return conn.requestTLS(
using: tlsConfiguration,
serverHostname: serverHostname,
logger: logger
).flatMapError { error in
conn.close().flatMapThrowing {
throw error
}
}.map { conn }
} else {
return eventLoop.makeSucceededFuture(conn)
}

let coders = PSQLConnection.Configuration.Coders(
jsonEncoder: PostgresJSONEncoderWrapper(_defaultJSONEncoder),
jsonDecoder: PostgresJSONDecoderWrapper(_defaultJSONDecoder)
)

let configuration = PSQLConnection.Configuration(
connection: .resolved(address: socketAddress, serverName: serverHostname),
authentication: nil,
tlsConfiguration: tlsConfiguration,
coders: coders)

return PSQLConnection.connect(
configuration: configuration,
logger: logger,
on: eventLoop
).map { connection in
PostgresConnection(underlying: connection, logger: logger)
}.flatMapErrorThrowing { error in
throw error.asAppropriatePostgresError
}
}
}


private final class PostgresErrorHandler: ChannelInboundHandler {
typealias InboundIn = Never

let logger: Logger
init(logger: Logger) {
self.logger = logger
}

func errorCaught(context: ChannelHandlerContext, error: Error) {
self.logger.error("Uncaught error: \(error)")
context.close(promise: nil)
context.fireErrorCaught(error)
}
}
Loading