diff --git a/src/api/API.jl b/src/api/API.jl index f4e0332..502ea16 100644 --- a/src/api/API.jl +++ b/src/api/API.jl @@ -229,14 +229,14 @@ _msgsizeof_parts(parts::Tuple) = msgsizeof(first(parts)) + _msgsizeof_parts(Base writepart(io, x) = write(io, x) function writepart(io, x::String) - occursin('\0', x) && throw(Postgres.PostgresInterfaceError("PostgreSQL protocol strings cannot contain a NUL byte")) + occursin('\0', x) && throw(PostgresInterfaceError("PostgreSQL protocol strings cannot contain a NUL byte")) write(io, x) write(io, UInt8(0)) end writepart(io, x::Integer) = write(io, hton(x)) function writepart(io, x::Tuple{String, String}) (occursin('\0', x[1]) || occursin('\0', x[2])) && - throw(Postgres.PostgresInterfaceError("PostgreSQL startup parameters cannot contain a NUL byte")) + throw(PostgresInterfaceError("PostgreSQL startup parameters cannot contain a NUL byte")) write(io, x[1]) write(io, UInt8(0)) write(io, x[2]) diff --git a/test/runtests.jl b/test/runtests.jl index 71a4f5a..c76008f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -983,6 +983,14 @@ include("notification_deadlines.jl") end end + @testset "Protocol Strings Reject NUL" begin + # the message is still buffered when the check throws, so nothing is sent + io = IOBuffer() + @test_throws Postgres.PostgresInterfaceError Postgres.API.writemessage(io, false, 'Q', "SELECT 1\0") + @test_throws Postgres.PostgresInterfaceError Postgres.API.writepart(io, ("user", "a\0b")) + @test position(io) == 0 + end + @testset "Seeded Protocol And DSN Fuzz" begin rng = MersenneTwister(0xf022) registry = Dict(Postgres.API.DEFAULT_TYPE_REGISTRY) @@ -1142,6 +1150,10 @@ include("notification_deadlines.jl") @test isopen(connp.socket) @test Tables.rowtable(DBInterface.execute(connp, "SELECT 1 AS a"))[1].a == 1 + # SQL with a NUL byte is rejected before anything is sent + @test_throws Postgres.PostgresInterfaceError DBInterface.execute(connp, "SELECT 1\0") + @test Tables.rowtable(DBInterface.execute(connp, "SELECT 1 AS a"))[1].a == 1 + # a consumer exception mid-result drains the remaining rows so # the connection stays usable stmt = DBInterface.prepare(connp, "SELECT i AS x, repeat('y', 10) AS s FROM generate_series(1, 200) i")