From 2b68038aa9f14d3c5570f04f22c22f46c29e34cb Mon Sep 17 00:00:00 2001 From: Jacob Quinn Date: Mon, 28 Sep 2026 14:01:42 -0600 Subject: [PATCH] Apply notification receive budgets to complete messages wait_for_notification(conn; timeout=...) stopped timing the read after the first byte of a PostgreSQL message, so a partly received message could block past the timeout. The timeout now covers the complete message, and partial PostgreSQL, TLS, and GSS input stays buffered for the next read. Idle validation and notification waits share the same complete-frame check. This is #16, which merged into the #15 branch after #15 had already landed, so it never reached main. It is ported onto current main with both test-registration lists kept and the docstring shortened. Co-Authored-By: Claude Opus 5.5 --- src/Postgres.jl | 115 ++++------- src/api/buffered.jl | 13 ++ test/gssapi_native_mock.jl | 1 + test/notification_deadlines.jl | 367 +++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 5 files changed, 429 insertions(+), 70 deletions(-) create mode 100644 test/notification_deadlines.jl diff --git a/src/Postgres.jl b/src/Postgres.jl index 6c56780..e4a8b81 100644 --- a/src/Postgres.jl +++ b/src/Postgres.jl @@ -234,16 +234,11 @@ function isvalid(conn::Connection) API.abort(socket) return false end - while bytesavailable(socket) >= 5 + while true message = try - pos = socket.pos - # Inspect framing without consuming an incomplete message. - len = Int(ntoh(reinterpret(Int32, socket.buffer[pos+1:pos+4])[1])) - 4 - (0 <= len <= API.MAX_MESSAGE_LEN) || throw(API.Error("invalid message length from server")) - bytesavailable(socket) >= 5 + len || break - frame = IOBuffer(@view(socket.buffer[pos:pos+4+len])) - socket.pos += 5 + len - mt, _ = API.readheader(frame, conn.debug) + frame = API.take_frame!(socket) + frame === nothing && break + mt, len = API.readheader(frame, conn.debug) if mt == UInt8('A') API.notificationResponse(len, frame) elseif mt == UInt8('N') @@ -541,91 +536,71 @@ end Block until a `NOTIFY` message arrives on the connection (see [`listen!`](@ref Postgres.listen!)) and return it as a [`Notification`](@ref Postgres.API.Notification). With a `timeout` (seconds), -return `nothing` if no message begins arriving in that window; once a message -starts, it is always read to completion so the connection is never left parked -mid-message. The connection lock is held while waiting, so use a dedicated -connection for listening — that is also the only way to receive every -notification, since notifications that arrive while the connection is busy -with a query are delivered to -[`notification_callback`](@ref Postgres.API.notification_callback) only during -the phases of a query that read result data. - -TLS and GSSAPI retain partial encrypted records across poll intervals. If the -timeout expires before a complete application byte is available, a later read -resumes that record. Once a PostgreSQL message begins, its remaining bytes are -read without a deadline. +return `nothing` if no complete notification arrives in time; a partly received +message stays buffered for the next call. Time spent waiting for the connection +lock, reconnecting, or running callbacks counts toward the timeout, but those +steps are not interrupted. + +The connection lock is held while waiting, so use a dedicated connection for +listening. Notifications received while a query uses the connection are instead +delivered to [`notification_callback`](@ref Postgres.API.notification_callback) +during the phases that read result data. """ function wait_for_notification(conn::Connection; timeout::Union{Real, Nothing}=nothing) - start_time = time() + start_time = time_ns() @lock conn.lock begin checkconn(conn) while true deadline_ns = if timeout === nothing Int64(time_ns()) + NOTIFICATION_POLL_INTERVAL_NS else - remaining_s = timeout - (time() - start_time) + remaining_s = timeout - (time_ns() - start_time) / 1_000_000_000 remaining_s <= 0 && return nothing # clamp before converting: an Inf or very large timeout would # overflow the nanosecond conversion remaining_ns = remaining_s >= 10.0 ? NOTIFICATION_POLL_INTERVAL_NS : round(Int64, remaining_s * 1_000_000_000) Int64(time_ns()) + min(NOTIFICATION_POLL_INTERVAL_NS, remaining_ns) end - # The deadline covers only the first byte: if it expires there, - # nothing of a message has been consumed and polling again is safe. - # Once a byte arrives the rest of the message is read without a - # deadline, so a message straddling the poll boundary can never - # leave the stream parked mid-message. - _set_read_deadline!(conn.socket, deadline_ns) - mt = try - read(conn.socket, UInt8) - catch err - if !_is_read_deadline_error(err) - # the stream position is unknowable, so the connection - # must never be reused. Transport deadlines are the - # exception: TLS and GSS retain partial record input. - close(conn.socket) - rethrow() - end - nothing - finally - # the deadline must be cleared on every path, including a - # rethrow: an expired deadline left set on the socket makes - # every later read on this connection fail - isopen(conn.socket) && _clear_read_deadline!(conn.socket) - end - # nothing arrived within this poll interval; nothing of a message - # has been consumed, so it is safe to loop and re-check the timeout - mt === nothing && continue - # A byte of a message has been consumed, so from here any failure - # leaves the stream at an unknowable position: close the connection - # rather than hand back one that still looks healthy. No deadline is - # in effect, so the message is read to completion. - # Read the message off the socket. Only the reading is guarded: - # once a message is fully consumed the stream is back at a clean - # boundary, so user callbacks and server errors are surfaced - # without destroying the connection. message = try - len = ntoh(read(conn.socket, Int32)) - 4 - (len < 0 || len > API.MAX_MESSAGE_LEN) && - throw(API.Error("invalid message length $len from server; connection protocol state is corrupted")) - conn.debug && @info "readheader: $(Char(mt)), $len" + frame = API.take_frame!(conn.socket) + if frame === nothing + conn.socket.ended && throw(EOFError()) + n = API.consume_input!(conn.socket) + if n === nothing + # Readiness is only a wake signal. Collect and parse + # nonblocking input again, even after raw TCP EOF, so + # TLS/GSS can diagnose truncated encrypted records. + transport = conn.socket.transport + tcp = transport isa Reseau.TCP.Conn ? transport : transport.tcp + _set_read_deadline!(conn.socket, deadline_ns) + try + eof(tcp) + finally + isopen(conn.socket) && _clear_read_deadline!(conn.socket) + end + end + continue + end + mt, len = API.readheader(frame, conn.debug) if mt == UInt8('A') - API.notificationResponse(len, conn.socket) + API.notificationResponse(len, frame) elseif mt == UInt8('N') - API.noticeResponse(len, conn.socket) + API.noticeResponse(len, frame) elseif mt == UInt8('S') - update_server_parameters!(conn, API.readbody(conn.socket, len)) + update_server_parameters!(conn, API.readbody(frame, len)) nothing elseif mt == UInt8('E') - API.errorResponse(len, conn.socket, conn.debug) + API.errorResponse(len, frame, conn.debug) else - API.skipbytes!(conn.socket, len) nothing end - catch - close(conn.socket) + catch err + _is_read_deadline_error(err) && continue + API.abort(conn.socket) rethrow() end + # Parsing has consumed a complete frame. User callback failures + # must propagate without discarding the aligned connection. if message isa API.Notification API.notification_callback(conn.style, message) return message @@ -633,7 +608,7 @@ function wait_for_notification(conn::Connection; timeout::Union{Real, Nothing}=n # FATAL and PANIC both terminate the session; close now so the # next use reports the error rather than a bare EOF from a # socket the server has already dropped - (message.severity == "FATAL" || message.severity == "PANIC") && close(conn.socket) + (message.severity == "FATAL" || message.severity == "PANIC") && API.abort(conn.socket) throw(message) elseif message !== nothing API.notice_callback(conn.style, message) diff --git a/src/api/buffered.jl b/src/api/buffered.jl index b1fcb9a..2c7d351 100644 --- a/src/api/buffered.jl +++ b/src/api/buffered.jl @@ -71,6 +71,19 @@ function consume_input!(io::BufferedConn)::Union{Int, Nothing} return n end +# Claim only a complete PostgreSQL message. Parse the returned view before +# collecting more input, which may compact or resize the shared buffer. +function take_frame!(io::BufferedConn) + bytesavailable(io) >= 5 || return nothing + pos = io.pos + len = Int(ntoh(reinterpret(Int32, @view(io.buffer[pos+1:pos+4]))[1])) - 4 + (0 <= len <= MAX_MESSAGE_LEN) || throw(Error("invalid message length from server")) + bytesavailable(io) >= 5 + len || return nothing + frame = IOBuffer(@view(io.buffer[pos:pos+4+len])) + io.pos += 5 + len + return frame +end + # On failure, abort TCP first: TLS close must not wait to send close_notify. function abort(io::BufferedConn) transport = io.transport diff --git a/test/gssapi_native_mock.jl b/test/gssapi_native_mock.jl index 7f21ed6..9f4f394 100644 --- a/test/gssapi_native_mock.jl +++ b/test/gssapi_native_mock.jl @@ -98,6 +98,7 @@ function test_gssapi() try Base.invokelatest(test_gssapi_protocol) Base.invokelatest(test_gss_fragmentation) + Base.invokelatest(test_notification_gss_deadlines) finally # Close even contexts left behind by a failed test before restoring the # real library; a finalizer must never hand a fake handle to Kerberos. diff --git a/test/notification_deadlines.jl b/test/notification_deadlines.jl new file mode 100644 index 0000000..81c85fd --- /dev/null +++ b/test/notification_deadlines.jl @@ -0,0 +1,367 @@ +# Scripted peers hold each fragment until the public wait has timed out. The +# watchdog only bounds cleanup if a reader ignores its receive deadline. +function notification_timeout(conn; timeout=0.02) + expired = Ref(false) + watchdog = Timer(5) do _ + expired[] = true + Postgres.API.abort(conn.socket) + end + try + @test Postgres.wait_for_notification(conn; timeout) === nothing + @test !expired[] + @test isopen(conn) + finally + close(watchdog) + end +end + +function with_notification_peer(f; gss=false) + peer = Channel{Any}(1) + handler = if gss + function (sock, _) + code, _ = read_request(sock) + code == GSSENC_REQUEST || error("expected GSS request") + send(sock, [UInt8('G')]) + io = server_handshake(sock) + read_startup(io) + ready!(io) + put!(peer, io) + drain(io) + end + else + serve_idle(sock -> (put!(peer, sock); drain(sock))) + end + with_fake_server(handler) do port, _ + conn = gss ? Postgres.Connection(host="127.0.0.1", port=port, user="u", dbname="db", + gssencmode="require", style=FakeGSSStyle()) : fake_connection(port, RecordingStyle()) + try + f(conn, take!(peer)) + finally + Postgres.API.abort(conn.socket) + close(conn) + end + end +end + +function notification_fragments(conn, send_fragment) + bytes = notification_msg(99, "deadline", "retained") + start = 1 + for stop in (1, 4, 5, length(bytes)-1) + send_fragment(bytes[start:stop]) + @test timedwait(() -> (Postgres.isvalid(conn); bytesavailable(conn.socket) == stop), 5) === :ok + notification_timeout(conn) + @test conn.socket.buffer[conn.socket.pos:end] == bytes[1:stop] + @test Postgres.isvalid(conn) + start = stop+1 + end + send_fragment(bytes[start:end]) + n = Postgres.wait_for_notification(conn; timeout=5) + @test n.channel == "deadline" + @test n.payload == "retained" + notification_timeout(conn) +end + +struct ThrowingNotificationStyle <: Postgres.API.AbstractPostgresStyle end +Postgres.API.notification_callback(::ThrowingNotificationStyle, n) = error("notification callback failed") + +function test_notification_deadlines() +@testset "Notification receive budget" begin + with_notification_peer() do conn, peer + notification_fragments(conn, bytes -> send(peer, bytes)) + @test length(conn.style.notifications) == 1 + # Side messages neither discard a partial notification nor change its + # receive budget. Unknown complete messages retain the skip policy. + send(peer, vcat(notice_msg("NOTICE", "00000", "before notification"), + param_status("application_name", "changed"), pgmsg('Y', UInt8[1, 2]), + notification_msg(1, "side", "next")[1:5])) + notification_timeout(conn) + @test length(conn.style.notices) == 1 + @test Postgres.get_server_parameter(conn, "application_name") == "changed" + send(peer, notification_msg(1, "side", "next")[6:end]) + @test Postgres.wait_for_notification(conn; timeout=5).payload == "next" + for timeout in (0, -1) + @test Postgres.wait_for_notification(conn; timeout) === nothing + end + for timeout in (nothing, Inf) + send(peer, notification_msg(1, "side", "unlimited")) + @test Postgres.wait_for_notification(conn; timeout).payload == "unlimited" + end + end + @testset "callback time reduces the remaining receive budget" begin + style = WaitingValidationStyle(Channel{Nothing}(1), Base.Event()) + with_fake_server(serve_idle(sock -> begin + send(sock, vcat(notice_msg("NOTICE", "00000", "wait"), + notification_msg(1, "after-callback", "buffered"))) + drain(sock) + end)) do port, _ + conn = fake_connection(port, style) + started = Ref(UInt64(0)) + reader = @async begin + started[] = time_ns() + Postgres.wait_for_notification(conn; timeout=1.0) + end + try + @test timedwait(() -> isready(style.entered) || istaskdone(reader), 5) === :ok + @test isready(style.entered) + timer = Timer(max(0.001, 1.02 - (time_ns()-started[])/1e9)) + try + wait(timer) + finally + close(timer) + end + @test !istaskdone(reader) # callbacks are deliberately not interrupted + notify(style.release) + @test fetch(reader) === nothing + @test Postgres.wait_for_notification(conn; timeout=5).payload == "buffered" + finally + notify(style.release) + close(conn) + wait(reader) + end + end + end + @testset "callback exceptions preserve aligned input" begin + for (style, bytes, text) in ( + (ThrowingValidationStyle(), notice_msg("NOTICE", "00000", "callback"), "validation callback failed"), + (ThrowingNotificationStyle(), notification_msg(1, "callback", "first"), "notification callback failed"), + ) + with_fake_server(serve_idle(sock -> begin + send(sock, vcat(bytes, param_status("application_name", "after-callback"))) + drain(sock) + end)) do port, _ + conn = fake_connection(port, style) + try + err = connect_err(() -> Postgres.wait_for_notification(conn; timeout=5)) + @test err isa ErrorException + @test occursin(text, sprint(showerror, err)) + @test isopen(conn) + notification_timeout(conn) + @test Postgres.get_server_parameter(conn, "application_name") == "after-callback" + finally + close(conn) + end + end + end + end + @testset "protocol and server errors" begin + for (bytes, then_close, error_type, remains_open) in ( + (vcat(UInt8('A'), be32(3)), false, Postgres.API.Error, false), + (notification_msg(1, "truncated", "body")[1:7], true, EOFError, false), + (UInt8['A', 0], true, EOFError, false), + (UInt8[], true, EOFError, false), + (error_response("ERROR", "XX000", "server error"), false, Postgres.API.Error, true), + (error_response("FATAL", "XX000", "server error"), false, Postgres.API.Error, false), + (error_response("PANIC", "XX000", "server error"), false, Postgres.API.Error, false), + ) + with_fake_server(serve_idle(sock -> begin + send(sock, bytes) + then_close ? close(sock) : drain(sock) + end)) do port, _ + conn = fake_connection(port, RecordingStyle()) + try + err = connect_err(() -> Postgres.wait_for_notification(conn; timeout=5)) + @test err isa error_type + @test isopen(conn) == remains_open + finally + close(conn) + end + end + end + end +end +end + +function test_notification_gss_deadlines() +@testset "Notification GSS receive budget" begin + with_notification_peer(gss=true) do conn, peer + notification_fragments(conn, bytes -> send(peer, bytes)) + payload = notification_msg(1, "gss-record", "resumed") + wire = vcat(be32(length(payload)+1), FAKE_SEAL, payload) + start = 1 + for stop in (1, 3, 4, length(wire)-1) + send(peer.sock, wire[start:stop]) + @test timedwait(() -> (Postgres.isvalid(conn); conn.socket.transport.received == stop), 5) === :ok + notification_timeout(conn) + @test conn.socket.transport.received == stop + start = stop+1 + end + send(peer.sock, wire[start:end]) + @test Postgres.wait_for_notification(conn; timeout=5).payload == "resumed" + end + for prefix in (0, 3, 7) + with_notification_peer(gss=true) do conn, peer + wire = vcat(be32(30), FAKE_SEAL, zeros(UInt8, 29)) + send(peer.sock, wire[1:prefix]) + close(peer.sock) + err = connect_err(() -> Postgres.wait_for_notification(conn; timeout=5)) + @test err isa (prefix == 0 ? EOFError : Postgres.API.Error) + @test !isopen(conn) + end + end +end + test_notification_gss_read_budget() +end + +const NotificationTransport = Union{Postgres.Reseau.TCP.Conn, Postgres.Reseau.TLS.Conn} +readn(io::NotificationTransport, n) = read!(io, Vector{UInt8}(undef, n)) +send(io::NotificationTransport, bytes) = write(io, bytes) + +function with_notification_tls_peer(f, config) + tcp = Postgres.Reseau.TCP + tls = Postgres.Reseau.TLS + listener = tcp.listen(tcp.loopback_addr(0)) + peer = Channel{tls.Conn}(1) + accepted = Ref{Union{Nothing, tcp.Conn}}(nothing) + task = @async begin + raw = tcp.accept(listener) + accepted[] = raw + try + code, _ = read_request(raw) + code == SSL_REQUEST || error("expected SSL request") + write(raw, UInt8('S')) + io = tls.server(raw, config) + tls.handshake!(io) + read_startup(io) + ready!(io) + put!(peer, io) + drain(io) + finally + close(raw) + end + end + watchdog = Timer(20) do _ + close(listener) + accepted[] === nothing || close(accepted[]) + end + conn = nothing + try + conn = Postgres.Connection(host="127.0.0.1", port=Int(tcp.addr(listener).port), + user="u", dbname="db", sslmode="require", gssencmode="disable", style=RecordingStyle()) + f(conn, take!(peer)) + finally + conn === nothing || Postgres.API.abort(conn.socket) + accepted[] === nothing || close(accepted[]) + close(listener) + close(watchdog) + wait(task) + end +end + +# Capture one record encrypted with this peer's live keys on another socket. +# Sending its bytes on the real connection permits splits inside TLS framing. +function notification_tls_record(peer, payload) + tcp = Postgres.Reseau.TCP + tls = Postgres.Reseau.TLS + listener = tcp.listen(tcp.loopback_addr(0)) + sender = tcp.connect(tcp.addr(listener)) + receiver = tcp.accept(listener) + try + state = peer.native_state + if state isa tls._TLS13NativeClientState + tls._tls13_write_record!(sender, state, tls._TLS_RECORD_TYPE_APPLICATION_DATA, payload) + else + tls._tls12_write_record!(sender, state, tls._TLS_RECORD_TYPE_APPLICATION_DATA, payload) + end + header = readn(receiver, 5) + return vcat(header, readn(receiver, (Int(header[4]) << 8) | Int(header[5]))) + finally + close(sender) + close(receiver) + close(listener) + end +end + +function test_notification_tls_deadlines() +@testset "Notification TLS receive budget" begin + if Sys.which("openssl") === nothing + @test_skip false + else + mktempdir() do dir + generate_ssl_material(dir) + tls = Postgres.Reseau.TLS + for version in (tls.TLS1_2_VERSION, tls.TLS1_3_VERSION) + config = tls.Config(verify_peer=false, cert_file=joinpath(dir, "server.crt"), + key_file=joinpath(dir, "server.key"), min_version=version, max_version=version) + with_notification_tls_peer(config) do conn, peer + notification_fragments(conn, bytes -> send(peer, bytes)) + wire = notification_tls_record(peer, notification_msg(1, "tls-record", "resumed")) + start = 1 + for stop in (1, 4, 5, length(wire)-1) + write(peer.tcp, wire[start:stop]) + @test timedwait(() -> (Postgres.isvalid(conn); conn.socket.transport.native_state.record_received == stop), 5) === :ok + notification_timeout(conn) + @test conn.socket.transport.native_state.record_received == stop + start = stop+1 + end + write(peer.tcp, wire[start:end]) + @test Postgres.wait_for_notification(conn; timeout=5).payload == "resumed" + @test length(conn.style.notifications) == 2 + end + for prefix in (0, 3, 7) + with_notification_tls_peer(config) do conn, peer + wire = notification_tls_record(peer, notification_msg(1, "tls-eof", "truncated")) + write(peer.tcp, wire[1:prefix]) + close(peer.tcp) + err = connect_err(() -> Postgres.wait_for_notification(conn; timeout=5)) + @test err isa (prefix == 0 ? EOFError : tls.TLSError) + @test !isopen(conn) + end + end + end + end + end +end +end + +function test_notification_gss_read_budget() + @testset "Notification in the sixteenth GSS record" begin + with_notification_peer(gss=true) do conn, peer + send(peer, notification_msg(1, "warmup", "complete")) + @test Postgres.wait_for_notification(conn; timeout=5).payload == "complete" + payload = notification_msg(1, "gss-budget", "sixteenth") + wire = vcat(repeat(vcat(be32(1), FAKE_SEAL), 15), be32(length(payload)+1), FAKE_SEAL, payload) + send(peer.sock, wire) + transport = conn.socket.transport + deadline = time_ns() + UInt64(5_000_000_000) + Postgres.API.set_read_deadline!(transport, deadline) + try + @test !eof(transport.tcp) + peeked = similar(wire) + socketops = Postgres.API.Reseau.SocketOps + while true + available = GC.@preserve peeked socketops.recv_from!( + transport.tcp.fd.pfd.sysfd, pointer(peeked), + Csize_t(length(peeked)), socketops.MSG_PEEK) + available > 0 || error("fixture socket lost readability") + available == length(peeked) && break + time_ns() < deadline || error("complete GSS fixture did not arrive") + yield() + end + @test peeked == wire + finally + Postgres.API.set_read_deadline!(transport, 0) + end + reader = @async Postgres.wait_for_notification(conn; timeout=5) + tcp = conn.socket.transport.tcp + waiter = Postgres.Reseau.IOPoll._poll_registration(tcp.fd.pfd.pd).read_waiter + watchdog = Timer(_ -> close(tcp), 5) + local n, completed_without_wait + try + while !istaskdone(reader) && (@atomic :acquire waiter.state) !== reader + yield() + end + completed_without_wait = istaskdone(reader) + n = fetch(reader) + finally + close(watchdog) + end + @test completed_without_wait # local plaintext must not cause a raw readiness wait + @test n !== nothing + if n !== nothing + @test n.channel == "gss-budget" + @test n.payload == "sixteenth" + end + notification_timeout(conn) + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 01bde22..e8d63f2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -499,6 +499,7 @@ include("gssapi.jl") include("isvalid_fake_server.jl") include("execute_fake_server.jl") include("executemany.jl") +include("notification_deadlines.jl") @testset "Postgres" begin test_timestamps() @@ -1022,6 +1023,8 @@ include("executemany.jl") test_isvalid_fragmentation() test_execute_fake_server() test_executemany_protocol() + test_notification_deadlines() + test_notification_tls_deadlines() require_integration = get(ENV, "POSTGRES_REQUIRE_INTEGRATION", "false") == "true" if !docker_available()