diff --git a/async_postgres/pg_protocol.nim b/async_postgres/pg_protocol.nim index 49c8df6..6dee54b 100644 --- a/async_postgres/pg_protocol.nim +++ b/async_postgres/pg_protocol.nim @@ -307,6 +307,16 @@ const ## tiny. Bounding by message length alone still buys a 16-byte `string` ## header per wire byte: 1 GiB in, 16 GiB preallocated, OutOfMemDefect. + MaxErrorOrNoticeFields* = 128 + ## Upper bound on fields in a single ErrorResponse/NoticeResponse. Only + ## about 20 field codes exist in the protocol; 128 leaves ample room for + ## future codes while capping worst-case allocation. + + MaxSaslMechanisms* = 64 + ## Upper bound on advertised SASL mechanisms in AuthenticationSASL. Only + ## SCRAM-SHA-256 and SCRAM-SHA-256-PLUS are standardised; 64 leaves room + ## for future mechanisms while bounding pre-auth allocation. + func makeBinarySafeLookup(): array[BinarySafeMaxOid + 1, bool] {.compileTime.} = for oid in BinarySafeOids: result[oid] = true @@ -816,6 +826,12 @@ proc parseAuthentication(body: openArray[byte]): BackendMessage = offset += consumed if mechanism.len == 0: break + # Bound advertised list before adding to survive a hostile pre-auth server. + if result.saslMechanisms.len >= MaxSaslMechanisms: + raise newException( + PgProtocolError, + "AuthenticationSASL: mechanism count exceeds maximum of " & $MaxSaslMechanisms, + ) result.saslMechanisms.add(mechanism) of 11: # SASLContinue @@ -877,6 +893,7 @@ proc parseErrorOrNotice(body: openArray[byte], isError: bool): BackendMessage = result = BackendMessage(kind: bmkNoticeResponse) result.noticeFields = @[] var offset = 0 + var fieldCount = 0 while offset < body.len: let fieldType = char(body[offset]) inc offset @@ -884,6 +901,15 @@ proc parseErrorOrNotice(body: openArray[byte], isError: bool): BackendMessage = break let (value, consumed) = decodeCString(body, offset) offset += consumed + # Bound field count before allocating: a hostile ~1 GiB body could otherwise + # produce hundreds of millions of two-byte fields, amplifying allocation. + if fieldCount >= MaxErrorOrNoticeFields: + let name = if isError: "ErrorResponse" else: "NoticeResponse" + raise newException( + PgProtocolError, + name & ": field count exceeds maximum of " & $MaxErrorOrNoticeFields, + ) + inc fieldCount let field = ErrorField(code: fieldType, value: value) if isError: result.errorFields.add(field) diff --git a/async_postgres/pg_replication.nim b/async_postgres/pg_replication.nim index bbd5621..d276761 100644 --- a/async_postgres/pg_replication.nim +++ b/async_postgres/pg_replication.nim @@ -586,6 +586,14 @@ proc identifySystem*( if results.len == 0 or results[0].rowCount == 0: raise newException(PgConnectionError, "IDENTIFY_SYSTEM returned no results") let qr = results[0] + # Guard fixed-column access so a malformed server response surfaces as a + # catchable PgConnectionError instead of an uncatchable IndexDefect from + # `cellInfo`. + if qr.fields.len < 3: + raise newException( + PgConnectionError, + "IDENTIFY_SYSTEM returned " & $qr.fields.len & " columns, expected >= 3", + ) let row = initRow(qr.data, 0) var info = SystemInfo() info.systemId = row.getStr(0) @@ -598,6 +606,11 @@ proc identifySystem*( return info proc decodeCreateSlotRow(qr: QueryResult): ReplicationSlotInfo = + if qr.fields.len < 2: + raise newException( + PgConnectionError, + "CREATE_REPLICATION_SLOT returned " & $qr.fields.len & " columns, expected >= 2", + ) let row = initRow(qr.data, 0) result.slotName = row.getStr(0) result.consistentPoint = parseLsn(row.getStr(1)) @@ -652,6 +665,11 @@ proc readReplicationSlot*( if results.len == 0 or results[0].rowCount == 0: raise newException(PgConnectionError, "READ_REPLICATION_SLOT returned no results") let qr = results[0] + if qr.fields.len < 2: + raise newException( + PgConnectionError, + "READ_REPLICATION_SLOT returned " & $qr.fields.len & " columns, expected >= 2", + ) let row = initRow(qr.data, 0) var info = ReplicationSlotInfo() # READ_REPLICATION_SLOT returns: slot_type, restart_lsn, restart_tli @@ -675,6 +693,11 @@ proc timelineHistory*( if results.len == 0 or results[0].rowCount == 0: raise newException(PgConnectionError, "TIMELINE_HISTORY returned no results") let qr = results[0] + if qr.fields.len < 2: + raise newException( + PgConnectionError, + "TIMELINE_HISTORY returned " & $qr.fields.len & " columns, expected >= 2", + ) let row = initRow(qr.data, 0) var info = TimelineHistory() if not row.isNull(0): diff --git a/tests/test_protocol.nim b/tests/test_protocol.nim index 7900209..f2225f2 100644 --- a/tests/test_protocol.nim +++ b/tests/test_protocol.nim @@ -1239,6 +1239,36 @@ suite "Backend decoding - edge cases": check res.state == psIncomplete check buf == original + test "ErrorResponse rejects excessive field count": + var body: seq[byte] = @[] + for _ in 0 .. MaxErrorOrNoticeFields: + body.add(byte('S')) + body.addCString("x") + body.add(0'u8) + var buf = buildMsg('E', body) + expect PgProtocolError: + discard parseBackendMessage(buf) + + test "NoticeResponse rejects excessive field count": + var body: seq[byte] = @[] + for _ in 0 .. MaxErrorOrNoticeFields: + body.add(byte('S')) + body.addCString("x") + body.add(0'u8) + var buf = buildMsg('N', body) + expect PgProtocolError: + discard parseBackendMessage(buf) + + test "AuthenticationSASL rejects excessive mechanism count": + var body: seq[byte] = @[] + body.addInt32(10) + for i in 0 .. MaxSaslMechanisms: + body.addCString("MECH-" & $i) + body.add(0'u8) + var buf = buildMsg('R', body) + expect PgProtocolError: + discard parseBackendMessage(buf) + suite "Frontend encoding - edge cases": test "encodeBind with result formats": let msg = encodeBind("", "", @[0'i16], @[some(@[byte('1')])], @[1'i16]) diff --git a/tests/test_replication.nim b/tests/test_replication.nim index b9a05c2..3bc3db9 100644 --- a/tests/test_replication.nim +++ b/tests/test_replication.nim @@ -781,3 +781,8 @@ suite "decodeCreateSlotRow": check info.consistentPoint == parseLsn("0/16B3740") check info.snapshotName == "" check info.outputPlugin == "" + + test "missing required columns raises catchable PgConnectionError": + let qr = mkSlotQr(["my_slot"], 1) + expect(PgConnectionError): + discard decodeCreateSlotRow(qr)