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()