Skip to content

Commit e9b90b2

Browse files
authored
Fix mishandling of SASL attribute parsing (#451)
1 parent 6ce96ab commit e9b90b2

File tree

2 files changed

+27
-4
lines changed

2 files changed

+27
-4
lines changed

Sources/PostgresNIO/Utilities/SASLAuthentication+SCRAM-SHA256.swift

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -209,14 +209,13 @@ fileprivate struct SCRAMMessageParser {
209209
}
210210

211211
static func parse(raw: [UInt8], isGS2Header: Bool = false) -> [SCRAMAttribute]? {
212-
213212
// There are two ways to implement this parse:
214213
// 1. All-at-once: Split on comma, split each on equals, validate
215214
// each results in a valid attribute.
216215
// 2. Sequential: State machine lookahead parse.
217216
// The former is simpler. The latter provides better validation.
218-
let likelyAttributeSets = raw.split(separator: .comma, maxSplits: isGS2Header ? 3 : Int.max, omittingEmptySubsequences: false)
219-
let likelyAttributePairs = likelyAttributeSets.map { $0.split(separator: .equals, maxSplits: 2, omittingEmptySubsequences: false) }
217+
let likelyAttributeSets = raw.split(separator: .comma, maxSplits: isGS2Header ? 2 : Int.max, omittingEmptySubsequences: false)
218+
let likelyAttributePairs = likelyAttributeSets.map { $0.split(separator: .equals, maxSplits: 1, omittingEmptySubsequences: false) }
220219

221220
let results = likelyAttributePairs.map { parseAttributePair(name: Array($0[0]), value: $0.dropFirst().first.map { Array($0) } ?? [], isGS2Header: isGS2Header) }
222221
let validResults = results.compactMap { $0 }
@@ -369,7 +368,7 @@ internal struct SHA256_PLUS: SASLAuthenticationMechanism {
369368
} // enum SCRAM
370369
} // enum SASLMechanism
371370

372-
/// Common impplementation of SCRAM-SHA-256 and SCRAM-SHA-256-PLUS
371+
/// Common implementation of SCRAM-SHA-256 and SCRAM-SHA-256-PLUS
373372
fileprivate final class SASLMechanism_SCRAM_SHA256_Common {
374373

375374
/// Initialized with initial client state

Tests/PostgresNIOTests/New/Connection State Machine/AuthenticationStateMachineTests.swift

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,30 @@ class AuthenticationStateMachineTests: XCTestCase {
4545
XCTAssertEqual(state.authenticationMessageReceived(.ok), .wait)
4646
}
4747

48+
func testAuthenticateSCRAMSHA256WithAtypicalEncoding() {
49+
let authContext = AuthContext(username: "test", password: "abc123", database: "test")
50+
var state = ConnectionStateMachine(requireBackendKeyData: true)
51+
XCTAssertEqual(state.connected(tls: .disable), .provideAuthenticationContext)
52+
XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext))
53+
54+
let saslResponse = state.authenticationMessageReceived(.sasl(names: ["SCRAM-SHA-256"]))
55+
guard case .sendSaslInitialResponse(name: let name, initialResponse: let responseData) = saslResponse else {
56+
return XCTFail("\(saslResponse) is not .sendSaslInitialResponse")
57+
}
58+
let responseString = String(decoding: responseData, as: UTF8.self)
59+
XCTAssertEqual(name, "SCRAM-SHA-256")
60+
XCTAssert(responseString.starts(with: "n,,n=test,r="))
61+
62+
let saslContinueResponse = state.authenticationMessageReceived(.saslContinue(data: .init(bytes:
63+
"r=\(responseString.dropFirst(12))RUJSZHhkeUVFNzRLNERKMkxmU05ITU1NZWcxaQ==,s=ijgUVaWgCDLRJyF963BKNA==,i=4096".utf8
64+
)))
65+
guard case .sendSaslResponse(let responseData2) = saslContinueResponse else {
66+
return XCTFail("\(saslContinueResponse) is not .sendSaslResponse")
67+
}
68+
let response2String = String(decoding: responseData2, as: UTF8.self)
69+
XCTAssertEqual(response2String.prefix(76), "c=biws,r=\(responseString.dropFirst(12))RUJSZHhkeUVFNzRLNERKMkxmU05ITU1NZWcxaQ==,p=")
70+
}
71+
4872
func testAuthenticationFailure() {
4973
let authContext = AuthContext(username: "test", password: "abc123", database: "test")
5074
var state = ConnectionStateMachine(requireBackendKeyData: true)

0 commit comments

Comments
 (0)