Skip to content

Commit 5280f31

Browse files
committed
psql authentication
1 parent 84f4881 commit 5280f31

20 files changed

+640
-309
lines changed

Package.swift

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,12 @@ let package = Package(
1212

1313
// Core extensions, type-aliases, and functions that facilitate common tasks.
1414
.package(url: "https://github.com/vapor/core.git", .branch("beta")),
15+
16+
// Non-blocking networking for Swift (HTTP and WebSockets).
17+
.package(url: "https://github.com/vapor/engine.git", .branch("beta")),
1518
],
1619
targets: [
17-
.target(name: "PostgreSQL", dependencies: ["Async", "Bits"]),
20+
.target(name: "PostgreSQL", dependencies: ["Async", "Bits", "TCP"]),
1821
.testTarget(name: "PostgreSQLTests", dependencies: ["PostgreSQL"]),
1922
]
2023
)
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import Bits
2+
3+
/// A frontend or backend PostgreSQL message.
4+
enum PostgreSQLMessage {
5+
case startupMessage(PostgreSQLStartupMessage)
6+
case errorResponse(PostgreSQLErrorResponse)
7+
case authenticationRequest(PostgreSQLAuthenticationRequest)
8+
case parameterStatus(PostgreSQLParameterStatus)
9+
case backendKeyData(PostgreSQLBackendKeyData)
10+
case readyForQuery(PostgreSQLReadyForQuery)
11+
}
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
import Foundation
2+
3+
/// Non-encoder wrapper for `_PostgreSQLMessageEncoder`.
4+
final class PostgreSQLMessageEncoder {
5+
/// Create a new `PostgreSQLMessageEncoder`
6+
init() {}
7+
8+
/// Encodes a `PostgreSQLMessage` to `Data`.
9+
func encode(_ message: PostgreSQLMessage) throws -> Data {
10+
let encoder = _PostgreSQLMessageEncoder()
11+
switch message {
12+
case .startupMessage(let message): try message.encode(to: encoder)
13+
default: fatalError("Unsupported encodable type")
14+
}
15+
encoder.updateSize()
16+
return encoder.data
17+
}
18+
}
19+
20+
// MARK: Encoder / Single
21+
22+
/// Internal `Encoder` implementation for the `PostgreSQLMessageEncoder`.
23+
internal final class _PostgreSQLMessageEncoder: Encoder, SingleValueEncodingContainer {
24+
/// See Encoder.codingPath
25+
var codingPath: [CodingKey]
26+
27+
/// See Encoder.userInfo
28+
var userInfo: [CodingUserInfoKey: Any]
29+
30+
/// The data currently being encoded
31+
var data: Data
32+
33+
/// Creates a new internal `_PostgreSQLMessageEncoder`
34+
init() {
35+
self.codingPath = []
36+
self.userInfo = [:]
37+
/// Start with 4 bytes for the int32 size chunk
38+
self.data = Data([0, 0, 0, 0])
39+
}
40+
41+
/// Updates the int32 size chunk in the data.
42+
func updateSize() {
43+
let size = numericCast(data.count) as Int32
44+
data.withUnsafeMutableBytes { (pointer: UnsafeMutablePointer<Int32>) in
45+
pointer.pointee = size.bigEndian
46+
}
47+
}
48+
49+
/// See Encoder.singleValueContainer
50+
func singleValueContainer() -> SingleValueEncodingContainer {
51+
return self
52+
}
53+
54+
/// See SingleValueEncodingContainer.encode
55+
func encode(_ value: String) throws {
56+
let stringData = Data(value.utf8)
57+
self.data.append(stringData + [0]) // c style string
58+
}
59+
60+
/// See SingleValueEncodingContainer.encode
61+
func encode(_ value: Int8) throws {
62+
self.data.append(numericCast(value))
63+
}
64+
65+
/// See SingleValueEncodingContainer.encode
66+
func encode(_ value: Int16) throws {
67+
var value = value.bigEndian
68+
withUnsafeBytes(of: &value) { buffer in
69+
let buffer = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self)
70+
self.data.append(buffer, count: 2)
71+
}
72+
}
73+
74+
/// See SingleValueEncodingContainer.encode
75+
func encode(_ value: Int32) throws {
76+
var value = value.bigEndian
77+
withUnsafeBytes(of: &value) { buffer in
78+
let buffer = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self)
79+
self.data.append(buffer, count: 4)
80+
}
81+
}
82+
83+
/// See SingleValueEncodingContainer.encode
84+
func encode(_ value: Int64) throws {
85+
var value = value.bigEndian
86+
withUnsafeBytes(of: &value) { buffer in
87+
let buffer = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self)
88+
self.data.append(buffer, count: 8)
89+
}
90+
}
91+
92+
/// See SingleValueEncodingContainer.encode
93+
func encode<T>(_ value: T) throws where T : Encodable {
94+
try value.encode(to: self)
95+
}
96+
97+
// Unsupported
98+
99+
func container<Key>(keyedBy type: Key.Type) -> KeyedEncodingContainer<Key> where Key : CodingKey { fatalError("Unsupported type: keyed container") }
100+
func unkeyedContainer() -> UnkeyedEncodingContainer { fatalError("Unsupported type: unkeyed container") }
101+
func encode(_ value: Int) throws { fatalError("Unsupported type: \(type(of: value))") }
102+
func encode(_ value: UInt) throws { fatalError("Unsupported type: \(type(of: value))") }
103+
func encode(_ value: UInt8) throws { fatalError("Unsupported type: \(type(of: value))") }
104+
func encode(_ value: UInt16) throws { fatalError("Unsupported type: \(type(of: value))") }
105+
func encode(_ value: UInt32) throws { fatalError("Unsupported type: \(type(of: value))") }
106+
func encode(_ value: UInt64) throws { fatalError("Unsupported type: \(type(of: value))") }
107+
func encode(_ value: Float) throws { fatalError("Unsupported type: \(type(of: value))") }
108+
func encode(_ value: Double) throws { fatalError("Unsupported type: \(type(of: value))") }
109+
func encode(_ value: Bool) throws { fatalError("Unsupported type: \(type(of: value))") }
110+
func encodeNil() throws { fatalError("Unsupported type: nil") }
111+
}
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import Async
2+
import Bits
3+
import Foundation
4+
5+
final class PostgreSQLMessageSerializer: ByteSerializer {
6+
typealias Input = PostgreSQLMessage
7+
typealias Output = ByteBuffer
8+
9+
var state: ByteSerializerState<PostgreSQLMessageSerializer>
10+
let buffer: MutableByteBuffer
11+
12+
init(bufferSize: Int = 4096) {
13+
buffer = MutableByteBuffer(start: .allocate(capacity: bufferSize), count: bufferSize)
14+
state = .init()
15+
}
16+
17+
func serialize(_ message: PostgreSQLMessage, state: Data?) throws -> ByteSerializerResult<PostgreSQLMessageSerializer> {
18+
if let state = state {
19+
return serialize(data: state)
20+
}
21+
22+
let data = try PostgreSQLMessageEncoder().encode(message)
23+
return serialize(data: data)
24+
}
25+
26+
func serialize(data: Data) -> ByteSerializerResult<PostgreSQLMessageSerializer> {
27+
print("serialize: \(data.hexDebug)")
28+
let count = data.copyBytes(to: buffer)
29+
let view = ByteBuffer(start: buffer.baseAddress, count: count)
30+
if data.count > count {
31+
return .incomplete(view, state: data[count..<data.count])
32+
} else {
33+
return .complete(view)
34+
}
35+
}
36+
37+
deinit {
38+
buffer.baseAddress?.deallocate(capacity: buffer.count)
39+
}
40+
}
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/// Represents [String: String] parameters encoded
2+
/// as a list of strings separated by null terminators
3+
/// and finished by a single null terminator.
4+
struct PostgreSQLParameters: Codable, ExpressibleByDictionaryLiteral {
5+
/// The internal parameter storage.
6+
var storage: [String: String]
7+
8+
/// See Encodable.encode
9+
public func encode(to encoder: Encoder) throws {
10+
var container = encoder.singleValueContainer()
11+
for (key, val) in storage {
12+
try container.encode(key)
13+
try container.encode(val)
14+
}
15+
try container.encode("")
16+
}
17+
18+
/// See ExpressibleByDictionaryLiteral.init
19+
init(dictionaryLiteral elements: (String, String)...) {
20+
var storage = [String: String]()
21+
for (key, val) in elements {
22+
storage[key] = val
23+
}
24+
self.storage = storage
25+
}
26+
}
27+
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,10 @@
1-
import Bits
2-
3-
/// A frontend or backend PostgreSQL message.
4-
protocol PostgreSQLMessage: Codable {
5-
/// The first byte of a message identifies the message type
6-
static var identifier: Byte? { get }
7-
}
8-
91
/// First message sent from the frontend during startup.
10-
struct PostgreSQLStartupMessage: PostgreSQLMessage {
2+
struct PostgreSQLStartupMessage: Encodable {
3+
/// Creates a `PostgreSQLStartupMessage` with "3.0" as the protocol version.
4+
static func versionThree(parameters: PostgreSQLParameters) -> PostgreSQLStartupMessage {
5+
return .init(protocolVersion: 196608, parameters: parameters)
6+
}
7+
118
/// The protocol version number. The most significant 16 bits are the major
129
/// version number (3 for the protocol described here). The least significant
1310
/// 16 bits are the minor version number (0 for the protocol described here).
@@ -25,38 +22,10 @@ struct PostgreSQLStartupMessage: PostgreSQLMessage {
2522
self.parameters = parameters
2623
}
2724

28-
/// Creates a `PostgreSQLStartupMessage` with "3.0" as the protocol version.
29-
static func versionThree(parameters: PostgreSQLParameters) -> PostgreSQLStartupMessage {
30-
return .init(protocolVersion: 196608, parameters: parameters)
31-
}
32-
33-
/// See `PostgreSQLMessage.identifier`
34-
static let identifier: Byte? = nil
35-
}
36-
37-
/// Represents [String: String] parameters encoded
38-
/// as a list of strings separated by null terminators
39-
/// and finished by a single null terminator.
40-
struct PostgreSQLParameters: Codable, ExpressibleByDictionaryLiteral {
41-
/// The internal parameter storage.
42-
var storage: [String: String]
43-
4425
/// See Encodable.encode
45-
public func encode(to encoder: Encoder) throws {
46-
var container = encoder.singleValueContainer()
47-
for (key, val) in storage {
48-
try container.encode(key)
49-
try container.encode(val)
50-
}
51-
try container.encode("")
52-
}
53-
54-
/// See ExpressibleByDictionaryLiteral.init
55-
init(dictionaryLiteral elements: (String, String)...) {
56-
var storage = [String: String]()
57-
for (key, val) in elements {
58-
storage[key] = val
59-
}
60-
self.storage = storage
26+
func encode(to encoder: Encoder) throws {
27+
var single = encoder.singleValueContainer()
28+
try single.encode(protocolVersion)
29+
try single.encode(parameters)
6130
}
6231
}
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/// Authentication request returned by the server.
2+
enum PostgreSQLAuthenticationRequest: Decodable {
3+
/// Specifies that the authentication was successful.
4+
case ok
5+
6+
/// Specifies that a clear-text password is required.
7+
case plaintext
8+
9+
/// See Decodable.init
10+
init(from decoder: Decoder) throws {
11+
let single = try decoder.singleValueContainer()
12+
let type = try single.decode(Int32.self)
13+
switch type {
14+
case 0: self = .ok
15+
case 3: self = .plaintext
16+
default: fatalError("Unsupported auth method: \(type)")
17+
}
18+
}
19+
}
20+
21+
extension PostgreSQLAuthenticationRequest: CustomStringConvertible {
22+
/// CustomStringConvertible.description
23+
var description: String {
24+
switch self {
25+
case .ok: return "none"
26+
case .plaintext: return "plaintext password required"
27+
}
28+
}
29+
}
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
/// Identifies the message as cancellation key data.
2+
/// The frontend must save these values if it wishes to be able to
3+
/// issue CancelRequest messages later.
4+
struct PostgreSQLBackendKeyData: Decodable {
5+
/// The process ID of this backend.
6+
var processID: Int32
7+
/// The secret key of this backend.
8+
var secretKey: Int32
9+
10+
/// See Decodable.decode
11+
init(from decoder: Decoder) throws {
12+
let single = try decoder.singleValueContainer()
13+
self.processID = try single.decode(Int32.self)
14+
self.secretKey = try single.decode(Int32.self)
15+
}
16+
}
17+
18+
extension PostgreSQLBackendKeyData: CustomStringConvertible {
19+
/// CustomStringConvertible.description
20+
var description: String {
21+
return "processID: \(processID), secretKey: \(secretKey)"
22+
}
23+
}
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/// First message sent from the frontend during startup.
2+
struct PostgreSQLErrorResponse: Decodable, Error {
3+
/// The error messages.
4+
var strings: [String]
5+
6+
/// Creates a new `PostgreSQLErrorResponse`.
7+
init(strings: [String]) {
8+
self.strings = strings
9+
}
10+
11+
/// See Decodable.init
12+
init(from decoder: Decoder) throws {
13+
let single = try decoder.singleValueContainer()
14+
var strings: [String] = []
15+
parse: while true {
16+
let byte = try single.decode(UInt8.self)
17+
switch byte {
18+
case 0: break parse
19+
default:
20+
let string = try single.decode(String.self)
21+
strings.append(string)
22+
}
23+
}
24+
self.init(strings: strings)
25+
}
26+
}

0 commit comments

Comments
 (0)