From d6b4ab8dd42ed8dea49cad335d56f531240f8128 Mon Sep 17 00:00:00 2001 From: Andrey Borodin Date: Fri, 23 Oct 2026 12:31:02 +0500 Subject: [PATCH v2 2/2] Add compression to the frontend/backend protocol. Negotiate Zstandard compression at connection startup, with a server GUC and an off/prefer/zstd libpq policy. Compress backend DataRow and CopyData messages and frontend CopyData during COPY FROM, leaving queries, parameters, and control messages in their ordinary representation. Use a PQcommMethods layer for backend output and a separate msg_buffer for decoded frontend input. Retain compression history between bounded segments, including segments of a large message. End backend frames at ReadyForQuery and frontend frames at CopyDone or CopyFail. Keep buffering thresholds internal and omit runtime codec switching and tuning options. Include protocol, COPY, pipeline, reset, and tracing coverage. Preserve the incoming decompressor while discarding data after a COPY error, even if ReadyForQuery has already been sent. --- doc/src/sgml/config.sgml | 24 + doc/src/sgml/libpq.sgml | 50 ++ doc/src/sgml/protocol.sgml | 75 +- meson.build | 1 + src/backend/commands/copyfromparse.c | 22 +- src/backend/libpq/Makefile | 1 + src/backend/libpq/meson.build | 1 + src/backend/libpq/pqcomm_compress.c | 679 +++++++++++++++++ src/backend/replication/walsender.c | 26 +- src/backend/tcop/backend_startup.c | 20 +- src/backend/tcop/postgres.c | 20 +- src/backend/utils/misc/guc_parameters.dat | 7 + src/backend/utils/misc/guc_tables.c | 8 + src/backend/utils/misc/postgresql.conf.sample | 3 + src/include/libpq/libpq.h | 14 + src/include/libpq/protocol.h | 5 +- src/include/libpq/protocol_compression.h | 33 + src/interfaces/libpq/Makefile | 9 +- src/interfaces/libpq/fe-auth.c | 10 +- src/interfaces/libpq/fe-compress.c | 435 +++++++++++ src/interfaces/libpq/fe-connect.c | 29 +- src/interfaces/libpq/fe-exec.c | 23 + src/interfaces/libpq/fe-misc.c | 5 +- src/interfaces/libpq/fe-protocol3.c | 86 ++- src/interfaces/libpq/fe-trace.c | 13 +- src/interfaces/libpq/libpq-int.h | 25 +- src/interfaces/libpq/meson.build | 2 + src/interfaces/libpq/nls.mk | 1 + src/interfaces/libpq/t/007_compression.pl | 232 ++++++ src/interfaces/libpq/test/libpq_testclient.c | 703 +++++++++++++++++- src/interfaces/libpq/test/meson.build | 8 +- src/tools/pgindent/typedefs.list | 2 + 32 files changed, 2533 insertions(+), 39 deletions(-) create mode 100644 src/backend/libpq/pqcomm_compress.c create mode 100644 src/include/libpq/protocol_compression.h create mode 100644 src/interfaces/libpq/fe-compress.c create mode 100644 src/interfaces/libpq/t/007_compression.pl diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml index f7db2995c3e..35ec564a8c6 100644 --- a/doc/src/sgml/config.sgml +++ b/doc/src/sgml/config.sgml @@ -963,6 +963,30 @@ include_dir 'conf.d' + + + protocol_compression (enum) + + protocol_compression configuration parameter + + + + + Controls whether new client connections may negotiate compression of + data-bearing protocol messages. The supported values are + off, the default, and zstd when + PostgreSQL was built with Zstandard. + Reloading this parameter affects new connections; existing + connections retain their negotiated setting. + + + Compression workspace is allocated on demand for each connection + that uses compression and can be retained between commands for reuse. + This is additional to the memory needed for ordinary protocol + messages. + + + diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml index b213a266cb2..f1ab5371a28 100644 --- a/doc/src/sgml/libpq.sgml +++ b/doc/src/sgml/libpq.sgml @@ -2234,6 +2234,40 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname + + compression + + + Enables compression of data-bearing protocol messages. + The supported values are off, the default, + prefer, and zstd. + prefer requests Zstandard compression and continues + without compression if the server rejects the protocol extension. + zstd requires compression and fails the connection + if the server rejects it. If libpq was built + without Zstandard, prefer is equivalent to + off, while zstd is rejected. + + + + Compression starts only after authentication. It currently applies + to DataRow and CopyData messages + sent by the server, and to CopyData sent by the + client during COPY FROM. Queries, bind parameters, + and control messages are not compressed. + + + + + Protocol data is compressed before SSL or GSS encryption. Observing + compressed lengths can reveal information when attacker-controlled + and secret values occur in the same compressed response. Enable this + option only when that risk is acceptable for the application. + + + + + krbsrvname @@ -7745,6 +7779,12 @@ void PQtrace(PGconn *conn, FILE *stream); . + + A received CompressedData message is followed by its + decoded messages, marked with direction B*. + The compressed payload itself is not printed. + + On Windows, if the libpq library and an application are @@ -8930,6 +8970,16 @@ myEventProc(PGEventId evtId, void *evtInfo, void *passThrough) + + + + PGCOMPRESSION + + PGCOMPRESSION behaves the same as the connection parameter. + + + diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml index 84af9372ebe..4918cf5b9ab 100644 --- a/doc/src/sgml/protocol.sgml +++ b/doc/src/sgml/protocol.sgml @@ -335,9 +335,14 @@ - - (No supported protocol extensions are currently defined.) - + _pq_.compression + zstd + Client and server + Enables Zstandard compression of server-to-client + DataRow and CopyData messages, + and client-to-server CopyData messages during + COPY FROM, after authentication. See + . @@ -4655,6 +4660,70 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;" + + CompressedData (F & B) + + + + Byte1('z') + + + Identifies the message as a flushed segment of compressed protocol + data. + + + + + + Int32 + + + Length of message contents in bytes, including self. + + + + + + Byten + + + A flushed segment of a Zstandard stream. The concatenated + uncompressed data from successive segments forms a stream of + ordinary protocol messages. A backend segment may contain + DataRow or CopyData messages. A + frontend segment contains bytes from at most one + CopyData message. In either direction, an ordinary + protocol message may span multiple compressed segments. + + + + + + + The compression context is retained between CompressedData messages. + Ordinary uncompressed messages may occur between them when the + uncompressed stream is at a message boundary, and do not change the + context. If an ordinary protocol message is incomplete, the next + message must be another CompressedData message. Each compressed payload + may produce at most 16 MiB of uncompressed data. Its compressed size + must not exceed ZSTD_compressBound for that amount + of input (16 MiB plus 64 KiB). Each segment ends with a complete flush, + so block-header overhead is bounded separately for each segment. + The Zstandard window size must not exceed 64 KiB. + A CompressedData message cannot contain the end of one Zstandard frame + and the beginning of another. After a backend frame ends, another + cannot begin before ReadyForQuery. After a frontend + frame ends, another cannot begin before CopyDone or + CopyFail. An unfinished backend frame must end before + ReadyForQuery, and an unfinished frontend frame must + end before CopyDone or CopyFail. + A CompressedData message that produces no output may carry the frame + epilogue. No CompressedData message may be sent before the first + ReadyForQuery of the connection. + + + + CopyData (F & B) diff --git a/meson.build b/meson.build index 7e95fd5941d..2831c21733d 100644 --- a/meson.build +++ b/meson.build @@ -3631,6 +3631,7 @@ libpq_deps += [ ldap_r, libintl, ssl, + zstd, ] libpq_oauth_deps += [ diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c index 06450535779..0fc6eee06cb 100644 --- a/src/backend/commands/copyfromparse.c +++ b/src/backend/commands/copyfromparse.c @@ -273,8 +273,10 @@ CopyGetData(CopyFromState cstate, void *databuf, int minread, int maxread) /* Try to receive another message */ int mtype; int maxmsglen; + bool message_body_read; readmessage: + message_body_read = false; HOLD_CANCEL_INTERRUPTS(); pq_startmsgread(); mtype = pq_getbyte(); @@ -282,6 +284,23 @@ CopyGetData(CopyFromState cstate, void *databuf, int minread, int maxread) ereport(ERROR, (errcode(ERRCODE_CONNECTION_FAILURE), errmsg("unexpected EOF on client connection with an open transaction"))); +#ifdef USE_ZSTD + pq_check_protocol_compression_message(mtype); + if (mtype == PqMsg_CompressedData) + { + mtype = pq_get_compressed_message(cstate->fe_msgbuf); + if (mtype == 0) + { + RESUME_CANCEL_INTERRUPTS(); + goto readmessage; + } + if (mtype == EOF) + ereport(ERROR, + (errcode(ERRCODE_CONNECTION_FAILURE), + errmsg("unexpected EOF on client connection with an open transaction"))); + message_body_read = true; + } +#endif /* Validate message type and set packet size limit */ switch (mtype) { @@ -303,7 +322,8 @@ CopyGetData(CopyFromState cstate, void *databuf, int minread, int maxread) break; } /* Now collect the message body */ - if (pq_getmessage(cstate->fe_msgbuf, maxmsglen)) + if (!message_body_read && + pq_getmessage(cstate->fe_msgbuf, maxmsglen)) ereport(ERROR, (errcode(ERRCODE_CONNECTION_FAILURE), errmsg("unexpected EOF on client connection with an open transaction"))); diff --git a/src/backend/libpq/Makefile b/src/backend/libpq/Makefile index 98eb2a8242d..2c58592fb90 100644 --- a/src/backend/libpq/Makefile +++ b/src/backend/libpq/Makefile @@ -26,6 +26,7 @@ OBJS = \ hba.o \ ifaddr.o \ pqcomm.o \ + pqcomm_compress.o \ pqformat.o \ pqmq.o \ pqsignal.o diff --git a/src/backend/libpq/meson.build b/src/backend/libpq/meson.build index 8571f652844..bb54a7cabbc 100644 --- a/src/backend/libpq/meson.build +++ b/src/backend/libpq/meson.build @@ -12,6 +12,7 @@ backend_sources += files( 'hba.c', 'ifaddr.c', 'pqcomm.c', + 'pqcomm_compress.c', 'pqformat.c', 'pqmq.c', 'pqsignal.c', diff --git a/src/backend/libpq/pqcomm_compress.c b/src/backend/libpq/pqcomm_compress.c new file mode 100644 index 00000000000..6d59cc70e8c --- /dev/null +++ b/src/backend/libpq/pqcomm_compress.c @@ -0,0 +1,679 @@ +/*------------------------------------------------------------------------- + * + * pqcomm_compress.c + * Compression layer for backend protocol I/O. + * + * Copyright (c) 2026, PostgreSQL Global Development Group + * + * src/backend/libpq/pqcomm_compress.c + * + *------------------------------------------------------------------------- + */ +#include "postgres.h" + +#include "libpq/libpq.h" +#include "libpq/protocol.h" +#include "libpq/protocol_compression.h" +#include "miscadmin.h" +#include "port/pg_bswap.h" +#include "utils/memutils.h" + +int protocol_compression = PROTOCOL_COMPRESSION_OFF; + +#ifdef USE_ZSTD +/* Internal batching policy, independent of the negotiated size limits. */ +#define PQ_COMPRESSION_INPUT_SIZE (32 * 1024) +#define PQ_COMPRESSION_BUFFER_KEEP_SIZE (1024 * 1024) + +static const PQcommMethods *PrevPQcommMethods; +static bool PqCompressBusy; +static bool PqCompressionNegotiated; +static bool PqCompressionStarted; +static bool PqCompressionActive; +static bool PqCompressionFrameStarted; +static size_t PqCompressionSmallBytes; +static StringInfoData PqCompressionInput; +static StringInfoData PqCompressionOutput; +static StringInfoData PqDecompressionInput; +static StringInfoData PqDecompressionBuffer; +static bool PqDecompressionBufferInitialized; +static bool PqDecompressionFrameStarted; +static bool PqDecompressionFrameEnded; +static ZSTD_CCtx *PqCompressionContext; +static ZSTD_DCtx *PqDecompressionContext; + +static void pq_compress_comm_reset(void); +static int pq_compress_flush(void); +static int pq_compress_flush_if_writable(void); +static bool pq_compress_is_send_pending(void); +static int pq_compress_putmessage(char msgtype, const char *s, size_t len); +static void pq_compress_putmessage_noblock(char msgtype, const char *s, size_t len); +static int pq_compress_putmessage_internal(char msgtype, const char *s, + size_t len, bool block); +static int pq_compress_flush_buffer(ZSTD_EndDirective directive, bool block); +static bool compression_buffer_init(StringInfo buf); +static bool compression_buffer_enlarge(StringInfo buf, size_t size); +static void compression_buffer_release(StringInfo buf); +static void pq_decompress_reset(void); + +static const PQcommMethods PqCommCompressMethods = { + .comm_reset = pq_compress_comm_reset, + .flush = pq_compress_flush, + .flush_if_writable = pq_compress_flush_if_writable, + .is_send_pending = pq_compress_is_send_pending, + .putmessage = pq_compress_putmessage, + .putmessage_noblock = pq_compress_putmessage_noblock +}; + +void +pq_enable_protocol_compression(void) +{ + Assert(!PqCompressionNegotiated); + PqCompressionNegotiated = true; + PrevPQcommMethods = PqCommMethods; + PqCommMethods = &PqCommCompressMethods; +} + +static bool +compression_buffer_init(StringInfo buf) +{ + buf->data = MemoryContextAllocExtended(TopMemoryContext, + STRINGINFO_DEFAULT_SIZE, + MCXT_ALLOC_NO_OOM); + if (buf->data == NULL) + return false; + buf->len = 0; + buf->maxlen = STRINGINFO_DEFAULT_SIZE; + buf->cursor = 0; + buf->data[0] = '\0'; + return true; +} + +static bool +compression_buffer_enlarge(StringInfo buf, size_t size) +{ + size_t needed; + size_t newlen; + char *data; + + if (size >= (size_t) MaxAllocSize - buf->len) + return false; + needed = (size_t) buf->len + size + 1; + if (needed <= (size_t) buf->maxlen) + return true; + + newlen = 2 * (size_t) buf->maxlen; + while (needed > newlen) + newlen = 2 * newlen; + if (!AllocSizeIsValid(newlen)) + newlen = needed; + + data = repalloc_extended(buf->data, newlen, MCXT_ALLOC_NO_OOM); + if (data == NULL) + return false; + buf->data = data; + buf->maxlen = (int) newlen; + return true; +} + +static void +compression_buffer_release(StringInfo buf) +{ + char *data; + + if (buf->maxlen <= PQ_COMPRESSION_BUFFER_KEEP_SIZE) + return; + Assert(buf->len == 0); + data = repalloc_extended(buf->data, STRINGINFO_DEFAULT_SIZE, + MCXT_ALLOC_NO_OOM); + if (data == NULL) + return; + buf->data = data; + buf->maxlen = STRINGINFO_DEFAULT_SIZE; + buf->cursor = 0; + buf->data[0] = '\0'; +} + +static void +pq_decompress_reset(void) +{ + if (PqDecompressionContext != NULL) + { + size_t result; + + result = ZSTD_DCtx_reset(PqDecompressionContext, + ZSTD_reset_session_only); + if (ZSTD_isError(result)) + { + ZSTD_freeDCtx(PqDecompressionContext); + PqDecompressionContext = NULL; + } + } + if (PqDecompressionBufferInitialized) + { + resetStringInfo(&PqDecompressionInput); + resetStringInfo(&PqDecompressionBuffer); + } + PqDecompressionFrameStarted = false; + PqDecompressionFrameEnded = false; +} + +void +pq_check_protocol_compression_message(int msgtype) +{ + if (msgtype == PqMsg_CompressedData) + { + if (!PqCompressionStarted) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("received compressed data before protocol initialization completed"))); + if (PqDecompressionFrameEnded) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("received data after the compressed protocol stream ended"))); + PqDecompressionFrameStarted = true; + } + else if (msgtype == PqMsg_CopyDone || msgtype == PqMsg_CopyFail) + { + if (PqDecompressionFrameStarted && !PqDecompressionFrameEnded) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol stream was not terminated before COPY ended"))); + pq_decompress_reset(); + compression_buffer_release(&PqDecompressionInput); + compression_buffer_release(&PqDecompressionBuffer); + } +} + +int +pq_get_compressed_message(StringInfo s) +{ + size_t result = 1; + + if (!PqCompressionNegotiated) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("received compressed data without negotiated compression"))); + + if (!PqDecompressionBufferInitialized) + { + if (!compression_buffer_init(&PqDecompressionInput) || + !compression_buffer_init(&PqDecompressionBuffer)) + ereport(FATAL, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while decompressing protocol data."))); + PqDecompressionBufferInitialized = true; + } + + if (PqDecompressionContext == NULL) + { + size_t rc; + + PqDecompressionContext = ZSTD_createDCtx(); + if (PqDecompressionContext == NULL) + ereport(FATAL, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while creating Zstandard decompression context."))); + rc = ZSTD_DCtx_setParameter(PqDecompressionContext, + ZSTD_d_windowLogMax, PQ_COMPRESSION_WINDOW_LOG); + if (ZSTD_isError(rc)) + ereport(FATAL, + (errcode(ERRCODE_INTERNAL_ERROR), + errmsg("could not configure Zstandard decompression context: %s", + ZSTD_getErrorName(rc)))); + } + + resetStringInfo(&PqDecompressionBuffer); + for (;;) + { + ZSTD_inBuffer input; + size_t segment_size = 0; + + if (pq_getmessage(&PqDecompressionInput, + PQ_COMPRESSION_MAX_WRAPPER_SIZE + 4)) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("invalid compressed protocol message"))); + if (PqDecompressionInput.len == 0) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol message is empty"))); + + input.src = PqDecompressionInput.data; + input.size = PqDecompressionInput.len; + input.pos = 0; + for (;;) + { + ZSTD_outBuffer out; + size_t old_input_pos = input.pos; + size_t output_size; + + output_size = PQ_COMPRESSION_MAX_SEGMENT_SIZE - segment_size; + output_size = output_size == 0 ? 1 : + Min(ZSTD_DStreamOutSize(), output_size); + if (!compression_buffer_enlarge(&PqDecompressionBuffer, output_size)) + ereport(FATAL, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while decompressing protocol data."))); + out.dst = PqDecompressionBuffer.data + PqDecompressionBuffer.len; + out.size = output_size; + out.pos = 0; + result = ZSTD_decompressStream(PqDecompressionContext, &out, &input); + if (ZSTD_isError(result)) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("invalid compressed protocol message"))); + if (input.pos == old_input_pos && out.pos == 0) + { + if (input.pos == input.size) + break; + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("invalid compressed protocol message"))); + } + PqDecompressionBuffer.len += out.pos; + PqDecompressionBuffer.data[PqDecompressionBuffer.len] = '\0'; + segment_size += out.pos; + if (segment_size > PQ_COMPRESSION_MAX_SEGMENT_SIZE) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol message is too large"))); + if (result == 0) + { + if (input.pos != input.size) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol message contains multiple frames"))); + break; + } + if (input.pos == input.size && out.pos < out.size) + break; + } + resetStringInfo(&PqDecompressionInput); + PqDecompressionFrameEnded = (result == 0); + + if (PqDecompressionBuffer.len == 0) + { + if (result == 0) + return 0; + goto read_next_wrapper; + } + if (PqDecompressionBuffer.len >= 5) + { + uint32 message_length; + + if (PqDecompressionBuffer.data[0] != PqMsg_CopyData) + goto invalid_contents; + memcpy(&message_length, PqDecompressionBuffer.data + 1, 4); + message_length = pg_ntoh32(message_length); + if (message_length < 4 || message_length > PG_INT32_MAX) + goto invalid_contents; + if (message_length + 1 < PqDecompressionBuffer.len) + goto invalid_contents; + if (message_length + 1 == PqDecompressionBuffer.len) + { + resetStringInfo(s); + if (!compression_buffer_enlarge(s, + PqDecompressionBuffer.len - 5)) + ereport(FATAL, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while decompressing protocol data."))); + memcpy(s->data, PqDecompressionBuffer.data + 5, + PqDecompressionBuffer.len - 5); + s->len = PqDecompressionBuffer.len - 5; + s->data[s->len] = '\0'; + s->cursor = 0; + resetStringInfo(&PqDecompressionBuffer); + return PqMsg_CopyData; + } + } + if (result == 0) + goto invalid_contents; + +read_next_wrapper: + pq_startmsgread(); + if (pq_getbyte() != PqMsg_CompressedData) + goto invalid_contents; + if (PqDecompressionFrameEnded) + goto invalid_contents; + } + +invalid_contents: + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol message contains invalid messages"))); + pg_unreachable(); +} + +static int +pq_compress_init(void) +{ + StringInfoData input; + StringInfoData output; + ZSTD_CCtx *cctx; + size_t result; + + Assert(PqCompressionNegotiated); + Assert(PqCompressionContext == NULL); + if (!compression_buffer_init(&input)) + goto oom; + if (!compression_buffer_init(&output)) + { + pfree(input.data); + goto oom; + } + cctx = ZSTD_createCCtx(); + if (cctx == NULL) + { + pfree(input.data); + pfree(output.data); + goto oom; + } + result = ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, PQ_COMPRESSION_WINDOW_LOG); + if (ZSTD_isError(result)) + { + ZSTD_freeCCtx(cctx); + pfree(input.data); + pfree(output.data); + ereport(COMMERROR, + (errcode(ERRCODE_INTERNAL_ERROR), + errmsg("could not configure Zstandard compression window: %s", + ZSTD_getErrorName(result)))); + goto fail; + } + PqCompressionInput = input; + PqCompressionOutput = output; + PqCompressionContext = cctx; + return 0; + +oom: + ereport(COMMERROR, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while initializing protocol compression."))); +fail: + ClientConnectionLost = 1; + InterruptPending = 1; + return EOF; +} + +static int +pq_compress_flush_buffer(ZSTD_EndDirective directive, bool block) +{ + size_t output_size; + ZSTD_inBuffer input; + ZSTD_outBuffer out; + size_t result; + + /* Do not create an empty frame for a result that was kept uncompressed. */ + if (!PqCompressionFrameStarted) + return 0; + + output_size = ZSTD_compressBound(PqCompressionInput.len); + resetStringInfo(&PqCompressionOutput); + if (!compression_buffer_enlarge(&PqCompressionOutput, output_size)) + { + ereport(COMMERROR, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while compressing protocol data."))); + goto fail; + } + input.src = PqCompressionInput.data; + input.size = PqCompressionInput.len; + input.pos = 0; + out.dst = PqCompressionOutput.data; + out.size = output_size; + out.pos = 0; + + while (input.pos < input.size) + { + result = ZSTD_compressStream2(PqCompressionContext, &out, &input, + ZSTD_e_continue); + if (ZSTD_isError(result)) + goto zstd_error; + if (out.pos == out.size && input.pos < input.size) + goto output_too_small; + } + do + { + result = ZSTD_compressStream2(PqCompressionContext, &out, &input, + directive); + if (ZSTD_isError(result)) + goto zstd_error; + if (out.pos == out.size && result != 0) + goto output_too_small; + } while (result != 0); + + if (out.pos > 0) + { + if (block) + { + if (PrevPQcommMethods->putmessage(PqMsg_CompressedData, + PqCompressionOutput.data, out.pos)) + goto fail; + } + else + { + PG_TRY(); + { + PrevPQcommMethods->putmessage_noblock(PqMsg_CompressedData, + PqCompressionOutput.data, out.pos); + } + PG_CATCH(); + { + /* The compressor has advanced; do not retry this input. */ + resetStringInfo(&PqCompressionInput); + ClientConnectionLost = 1; + InterruptPending = 1; + PG_RE_THROW(); + } + PG_END_TRY(); + } + } + + resetStringInfo(&PqCompressionInput); + return 0; + +zstd_error: + ereport(COMMERROR, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("Zstandard compression failed: %s", + ZSTD_getErrorName(result)))); + goto fail; + +output_too_small: + ereport(COMMERROR, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("Zstandard compression output buffer is too small"))); + +fail: + + /* + * The compressor has consumed part of the pending input, so the segment + * we were building cannot be produced again. Drop it and give up on the + * connection, as the send path does for a socket error. + */ + resetStringInfo(&PqCompressionInput); + ClientConnectionLost = 1; + InterruptPending = 1; + return EOF; +} + +static int +pq_compress_append(const void *data, size_t len, bool block) +{ + const char *ptr = data; + + while (len > 0) + { + size_t available = PQ_COMPRESSION_INPUT_SIZE - + PqCompressionInput.len; + size_t part = Min(len, available); + + if (part == 0) + { + if (pq_compress_flush_buffer(ZSTD_e_flush, block)) + return EOF; + continue; + } + if (!compression_buffer_enlarge(&PqCompressionInput, part)) + { + ereport(COMMERROR, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"), + errdetail("Failed while compressing protocol data."))); + ClientConnectionLost = 1; + InterruptPending = 1; + return EOF; + } + appendBinaryStringInfo(&PqCompressionInput, ptr, part); + ptr += part; + len -= part; + if (PqCompressionInput.len == PQ_COMPRESSION_INPUT_SIZE) + { + if (pq_compress_flush_buffer(ZSTD_e_flush, block)) + return EOF; + /* Do not let a good compression ratio delay network delivery. */ + if (block && PrevPQcommMethods->flush()) + return EOF; + } + } + return 0; +} + +/* + * Ordinary messages remain visible. Compression starts only after the first + * ReadyForQuery, and every subsequent ReadyForQuery ends the current frame. + */ +static int +pq_compress_putmessage_internal(char msgtype, const char *s, size_t len, + bool block) +{ + uint32 n32; + + if (PqCompressionStarted && + (msgtype == PqMsg_DataRow || msgtype == PqMsg_CopyData)) + { + /* Keep isolated short rows out of the compressor. */ + if (!PqCompressionActive && len + 4 < 60) + { + PqCompressionSmallBytes += len + 5; + if (PqCompressionSmallBytes >= 1024) + PqCompressionActive = true; + } + else + { + if (PqCompressionContext == NULL && pq_compress_init()) + return EOF; + PqCompressionActive = true; + PqCompressionFrameStarted = true; + n32 = pg_hton32((uint32) (len + 4)); + if (pq_compress_append(&msgtype, 1, block) || + pq_compress_append(&n32, 4, block) || + pq_compress_append(s, len, block)) + return EOF; + return 0; + } + } + + if (PqCompressionStarted && + pq_compress_flush_buffer(msgtype == PqMsg_ReadyForQuery ? + ZSTD_e_end : ZSTD_e_flush, block)) + return EOF; + + if (block) + { + if (PrevPQcommMethods->putmessage(msgtype, s, len)) + return EOF; + } + else + PrevPQcommMethods->putmessage_noblock(msgtype, s, len); + + if (msgtype == PqMsg_ReadyForQuery) + { + PqCompressionStarted = true; + PqCompressionActive = false; + PqCompressionFrameStarted = false; + PqCompressionSmallBytes = 0; + + /* + * Do not reset the incoming stream here. After a COPY error, the + * frontend may still be sending data that we must decode and discard + * until CopyDone or CopyFail. + */ + } + return 0; +} + +static void +pq_compress_comm_reset(void) +{ + PqCompressBusy = false; + PrevPQcommMethods->comm_reset(); +} + +static int +pq_compress_flush(void) +{ + int result; + + if (PqCompressBusy) + return 0; + PqCompressBusy = true; + result = pq_compress_flush_buffer(ZSTD_e_flush, true); + if (result == 0) + result = PrevPQcommMethods->flush(); + PqCompressBusy = false; + return result; +} + +static int +pq_compress_flush_if_writable(void) +{ + int result; + + if (PqCompressBusy) + return 0; + PqCompressBusy = true; + result = pq_compress_flush_buffer(ZSTD_e_flush, false); + if (result == 0) + result = PrevPQcommMethods->flush_if_writable(); + PqCompressBusy = false; + return result; +} + +static bool +pq_compress_is_send_pending(void) +{ + return PqCompressionInput.len > 0 || + PrevPQcommMethods->is_send_pending(); +} + +static int +pq_compress_putmessage(char msgtype, const char *s, size_t len) +{ + int result; + + if (PqCompressBusy) + return 0; + PqCompressBusy = true; + result = pq_compress_putmessage_internal(msgtype, s, len, true); + PqCompressBusy = false; + return result; +} + +static void +pq_compress_putmessage_noblock(char msgtype, const char *s, size_t len) +{ + if (PqCompressBusy) + return; + PqCompressBusy = true; + (void) pq_compress_putmessage_internal(msgtype, s, len, false); + PqCompressBusy = false; +} +#endif diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c index 9f21c7c8dfe..06783c40b57 100644 --- a/src/backend/replication/walsender.c +++ b/src/backend/replication/walsender.c @@ -789,9 +789,13 @@ HandleUploadManifestPacket(StringInfo buf, off_t *offset, { int mtype; int maxmsglen; + bool message_body_read; +#ifdef USE_ZSTD +read_message: +#endif + message_body_read = false; HOLD_CANCEL_INTERRUPTS(); - pq_startmsgread(); mtype = pq_getbyte(); if (mtype == EOF) @@ -799,6 +803,24 @@ HandleUploadManifestPacket(StringInfo buf, off_t *offset, (errcode(ERRCODE_CONNECTION_FAILURE), errmsg("unexpected EOF on client connection with an open transaction"))); +#ifdef USE_ZSTD + pq_check_protocol_compression_message(mtype); + if (mtype == PqMsg_CompressedData) + { + mtype = pq_get_compressed_message(buf); + if (mtype == 0) + { + RESUME_CANCEL_INTERRUPTS(); + goto read_message; + } + if (mtype == EOF) + ereport(ERROR, + (errcode(ERRCODE_CONNECTION_FAILURE), + errmsg("unexpected EOF on client connection with an open transaction"))); + message_body_read = true; + } +#endif + switch (mtype) { case PqMsg_CopyData: @@ -820,7 +842,7 @@ HandleUploadManifestPacket(StringInfo buf, off_t *offset, } /* Now collect the message body */ - if (pq_getmessage(buf, maxmsglen)) + if (!message_body_read && pq_getmessage(buf, maxmsglen)) ereport(ERROR, (errcode(ERRCODE_CONNECTION_FAILURE), errmsg("unexpected EOF on client connection with an open transaction"))); diff --git a/src/backend/tcop/backend_startup.c b/src/backend/tcop/backend_startup.c index 912ad7dc957..740dcdf4f39 100644 --- a/src/backend/tcop/backend_startup.c +++ b/src/backend/tcop/backend_startup.c @@ -759,6 +759,7 @@ retry: { int32 offset = sizeof(ProtocolVersion); List *unrecognized_protocol_options = NIL; + bool compression_requested = false; /* * Scan packet body for name/option pairs. We can assume any string @@ -808,11 +809,28 @@ retry: valptr), errhint("Valid values are: \"false\", 0, \"true\", 1, \"database\"."))); } + else if (strcmp(nameptr, "_pq_.compression") == 0) + { + if (compression_requested) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("protocol compression option specified more than once"))); + compression_requested = true; +#ifdef USE_ZSTD + if (strcmp(valptr, "zstd") == 0 && + protocol_compression == PROTOCOL_COMPRESSION_ZSTD) + pq_enable_protocol_compression(); + else +#endif + /* Reject unknown future profiles through negotiation. */ + unrecognized_protocol_options = + lappend(unrecognized_protocol_options, pstrdup(nameptr)); + } else if (strncmp(nameptr, "_pq_.", 5) == 0) { /* * Any option beginning with _pq_. is reserved for use as a - * protocol-level option, but at present no such options are + * protocol-level option, but at present only one option is * defined. */ unrecognized_protocol_options = diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c index b6bdfe213fe..92a6387d8da 100644 --- a/src/backend/tcop/postgres.c +++ b/src/backend/tcop/postgres.c @@ -367,11 +367,16 @@ SocketBackend(StringInfo inBuf) { int qtype; int maxmsglen; + bool message_body_read = false; /* * Get message type code from the frontend. */ HOLD_CANCEL_INTERRUPTS(); + +#ifdef USE_ZSTD +read_message: +#endif pq_startmsgread(); qtype = pq_getbyte(); @@ -396,6 +401,19 @@ SocketBackend(StringInfo inBuf) return qtype; } +#ifdef USE_ZSTD + pq_check_protocol_compression_message(qtype); + if (qtype == PqMsg_CompressedData) + { + qtype = pq_get_compressed_message(inBuf); + if (qtype == 0) + goto read_message; + if (qtype == EOF) + return EOF; + message_body_read = true; + } +#endif + /* * Validate message type code before trying to read body; if we have lost * sync, better to say "command unknown" than to run out of memory because @@ -476,7 +494,7 @@ SocketBackend(StringInfo inBuf) * after the type code; we can read the message contents independently of * the type. */ - if (pq_getmessage(inBuf, maxmsglen)) + if (!message_body_read && pq_getmessage(inBuf, maxmsglen)) return EOF; /* suitable message already logged */ RESUME_CANCEL_INTERRUPTS(); diff --git a/src/backend/utils/misc/guc_parameters.dat b/src/backend/utils/misc/guc_parameters.dat index 380679a76c0..c871aa74258 100644 --- a/src/backend/utils/misc/guc_parameters.dat +++ b/src/backend/utils/misc/guc_parameters.dat @@ -2423,6 +2423,13 @@ check_hook => 'check_primary_slot_name', }, +{ name => 'protocol_compression', type => 'enum', context => 'PGC_SIGHUP', group => 'CONN_AUTH_SETTINGS', + short_desc => 'Sets the compression method allowed for new client connections.', + variable => 'protocol_compression', + boot_val => 'PROTOCOL_COMPRESSION_OFF', + options => 'protocol_compression_options', +}, + { name => 'quote_all_identifiers', type => 'bool', context => 'PGC_USERSET', group => 'COMPAT_OPTIONS_PREVIOUS', short_desc => 'When generating SQL fragments, quote all identifiers.', variable => 'quote_all_identifiers', diff --git a/src/backend/utils/misc/guc_tables.c b/src/backend/utils/misc/guc_tables.c index 342aaeef59a..63964ff1a64 100644 --- a/src/backend/utils/misc/guc_tables.c +++ b/src/backend/utils/misc/guc_tables.c @@ -477,6 +477,14 @@ static const struct config_enum_entry default_toast_compression_options[] = { {NULL, 0, false} }; +static const struct config_enum_entry protocol_compression_options[] = { + {"off", PROTOCOL_COMPRESSION_OFF, false}, +#ifdef USE_ZSTD + {"zstd", PROTOCOL_COMPRESSION_ZSTD, false}, +#endif + {NULL, 0, false} +}; + static const struct config_enum_entry wal_compression_options[] = { {"pglz", WAL_COMPRESSION_PGLZ, false}, #ifdef USE_LZ4 diff --git a/src/backend/utils/misc/postgresql.conf.sample b/src/backend/utils/misc/postgresql.conf.sample index 01f98de8a9f..7291f5e2eba 100644 --- a/src/backend/utils/misc/postgresql.conf.sample +++ b/src/backend/utils/misc/postgresql.conf.sample @@ -826,6 +826,9 @@ #createrole_self_grant = '' # set and/or inherit #event_triggers = on +# - Protocol Compression - +#protocol_compression = off # off, zstd + # - Locale and Formatting - #datestyle = 'iso, mdy' diff --git a/src/include/libpq/libpq.h b/src/include/libpq/libpq.h index d15073a0a93..c52eef8d88f 100644 --- a/src/include/libpq/libpq.h +++ b/src/include/libpq/libpq.h @@ -175,4 +175,18 @@ extern bool check_ssl_key_file_permissions(const char *ssl_key_file, bool isServerStart); extern HostsFileLoadResult load_hosts(List **hosts, char **err_msg); +/* Protocol compression is negotiated at connection startup. */ +typedef enum ProtocolCompression +{ + PROTOCOL_COMPRESSION_OFF, + PROTOCOL_COMPRESSION_ZSTD +} ProtocolCompression; + +extern PGDLLIMPORT int protocol_compression; +#ifdef USE_ZSTD +extern void pq_enable_protocol_compression(void); +extern void pq_check_protocol_compression_message(int msgtype); +extern int pq_get_compressed_message(StringInfo s); +#endif + #endif /* LIBPQ_H */ diff --git a/src/include/libpq/protocol.h b/src/include/libpq/protocol.h index 503852ca988..885bd24930a 100644 --- a/src/include/libpq/protocol.h +++ b/src/include/libpq/protocol.h @@ -61,8 +61,9 @@ /* These are the codes sent by both the frontend and backend. */ -#define PqMsg_CopyDone 'c' -#define PqMsg_CopyData 'd' +#define PqMsg_CopyDone 'c' +#define PqMsg_CopyData 'd' +#define PqMsg_CompressedData 'z' /* Additional codes sent by parallel workers to leader processes. */ diff --git a/src/include/libpq/protocol_compression.h b/src/include/libpq/protocol_compression.h new file mode 100644 index 00000000000..0ae10a73d77 --- /dev/null +++ b/src/include/libpq/protocol_compression.h @@ -0,0 +1,33 @@ +/*------------------------------------------------------------------------- + * + * protocol_compression.h + * Size limits for compressed frontend/backend protocol messages. + * + * Copyright (c) 2026, PostgreSQL Global Development Group + * + * src/include/libpq/protocol_compression.h + * + *------------------------------------------------------------------------- + */ +#ifndef PROTOCOL_COMPRESSION_H +#define PROTOCOL_COMPRESSION_H + +#ifdef USE_ZSTD +#include +#endif + +/* Limits of the negotiated profile, not buffer allocation targets. */ +#define PQ_COMPRESSION_MAX_SEGMENT_SIZE (16 * 1024 * 1024) +#define PQ_COMPRESSION_WINDOW_LOG 16 + +#ifdef USE_ZSTD +/* + * Bound a single segment, not the entire stream. Each segment starts after + * the preceding flush has completed and has no intermediate flushes. + * Otherwise, repeated small blocks could exceed the single-pass bound. + */ +#define PQ_COMPRESSION_MAX_WRAPPER_SIZE \ + ZSTD_COMPRESSBOUND(PQ_COMPRESSION_MAX_SEGMENT_SIZE) +#endif + +#endif /* PROTOCOL_COMPRESSION_H */ diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 0963995eed4..16e0bcdbf04 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -33,6 +33,7 @@ OBJS = \ $(WIN32RES) \ fe-auth-scram.o \ fe-cancel.o \ + fe-compress.o \ fe-connect.o \ fe-exec.o \ fe-lobj.o \ @@ -87,9 +88,9 @@ endif # that are built correctly for use in a shlib. SHLIB_LINK_INTERNAL = -lpgcommon_shlib -lpgport_shlib ifneq ($(PORTNAME), win32) -SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -lsocket -lnsl -lresolv -lintl -ldl -lm, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS) +SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -lsocket -lnsl -lresolv -lintl -ldl -lm -lzstd, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS) else -SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi32 -lssl -lsocket -lnsl -lresolv -lintl -lm $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE) +SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi32 -lssl -lsocket -lnsl -lresolv -lintl -lm -lzstd $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE) endif ifeq ($(PORTNAME), win32) SHLIB_LINK += -lshell32 -lws2_32 -lsecur32 $(filter -lcomerr32 -lkrb5_32, $(LIBS)) @@ -109,6 +110,10 @@ $(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,libssl) $(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,libcrypto) endif +ifneq (,$(filter -lzstd, $(LIBS))) +$(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,libzstd) +endif + ifeq ($(with_libcurl),yes) # libpq.so doesn't link against libcurl, but libpq.a needs libpq-oauth, and # libpq-oauth needs libcurl. Put both into *.private. diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c index 8df5bad92ab..6cd43929699 100644 --- a/src/interfaces/libpq/fe-auth.c +++ b/src/interfaces/libpq/fe-auth.c @@ -99,7 +99,7 @@ pg_GSS_continue(PGconn *conn, int payloadlen) /* finished parsing, trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false); + pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false, false); /* Only try to acquire credentials if GSS delegation isn't disabled. */ if (!pg_GSS_have_cred_cache(&conn->gcred)) @@ -267,7 +267,7 @@ pg_SSPI_continue(PGconn *conn, int payloadlen) /* finished parsing, trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false); + pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false, false); OutBuffers[0].pvBuffer = NULL; OutBuffers[0].BufferType = SECBUFFER_TOKEN; @@ -605,7 +605,7 @@ pg_SASL_init(PGconn *conn, int payloadlen, bool *async) /* finished parsing, trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false); + pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false, false); Assert(conn->sasl); @@ -726,7 +726,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final, bool *async) /* finished parsing, trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false); + pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false, false); /* For safety and convenience, ensure the buffer is NULL-terminated. */ challenge[payloadlen] = '\0'; @@ -809,7 +809,7 @@ pg_password_sendauth(PGconn *conn, const char *password, AuthRequest areq) /* finished parsing, trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false); + pqTraceOutputMessage(conn, conn->inBuffer.buffer + conn->inBuffer.start, false, false); /* Encrypt the password if needed. */ diff --git a/src/interfaces/libpq/fe-compress.c b/src/interfaces/libpq/fe-compress.c new file mode 100644 index 00000000000..aa6719fece8 --- /dev/null +++ b/src/interfaces/libpq/fe-compress.c @@ -0,0 +1,435 @@ +/*------------------------------------------------------------------------- + * + * fe-compress.c + * Compression support for the frontend/backend protocol. + * + * Portions Copyright (c) 2026, PostgreSQL Global Development Group + * + * src/interfaces/libpq/fe-compress.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include + +#include "libpq-fe.h" +#include "libpq-int.h" +#include "libpq/protocol_compression.h" +#include "port/pg_bswap.h" + +#define PQ_COMPRESSION_BUFFER_KEEP_SIZE (1024 * 1024) + +void +pqCompressionReset(PGconn *conn) +{ +#ifdef USE_ZSTD + if (conn->compression_dctx != NULL) + { + ZSTD_freeDCtx((ZSTD_DCtx *) conn->compression_dctx); + conn->compression_dctx = NULL; + } + if (conn->compression_cctx != NULL) + { + ZSTD_freeCCtx((ZSTD_CCtx *) conn->compression_cctx); + conn->compression_cctx = NULL; + } +#endif + if (conn->compression_buffers_initialized) + { + termPQExpBuffer(&conn->compression_output_buffer); + conn->compression_buffers_initialized = false; + } + free(conn->decompressBuffer.buffer); + memset(&conn->decompressBuffer, 0, sizeof(msg_buffer)); + conn->compression_in_frame = false; + conn->compression_frame_ended = false; + conn->compression_ready = false; + conn->compression_copy_started = false; + conn->compression_rejected = false; +} + +/* + * Choose between network input and decoded messages. Only a complete inner + * message is handed to the existing parser; an incomplete one stays here + * while the next wrapper is read. Never copy decoded messages into inBuffer. + */ +msg_buffer * +pqGetMessageBuffer(PGconn *conn) +{ + msg_buffer *buf = &conn->decompressBuffer; + int available = buf->end - buf->start; + + if (available >= 5) + { + uint32 len; + + memcpy(&len, buf->buffer + buf->start + 1, 4); + len = pg_ntoh32(len); + if ((buf->buffer[buf->start] != PqMsg_DataRow && + buf->buffer[buf->start] != PqMsg_CopyData) || + len < 4 || len >= INT_MAX) + goto invalid; + if ((size_t) available >= (size_t) len + 1) + return buf; + } + if (available && conn->compression_frame_ended) + goto invalid; + + buf = &conn->inBuffer; + if (buf->end - buf->start >= 5) + { + char id = buf->buffer[buf->start]; + uint32 len; + + memcpy(&len, buf->buffer + buf->start + 1, 4); + len = pg_ntoh32(len); + if (id == PqMsg_CompressedData) + { + if (!conn->compression_ready || len <= 4) + goto invalid; +#ifdef USE_ZSTD + if (len - 4 > PQ_COMPRESSION_MAX_WRAPPER_SIZE) + goto invalid; +#endif + } + else if (available || + (id == PqMsg_ReadyForQuery && conn->compression_in_frame)) + goto invalid; + } + return buf; + +invalid: + libpq_append_conn_error(conn, "invalid compressed protocol message"); + return NULL; +} + +void +pqCompressionReady(PGconn *conn) +{ +#ifdef USE_ZSTD + char *data; + + if (conn->compression_copy_started) + { + ZSTD_freeCCtx((ZSTD_CCtx *) conn->compression_cctx); + conn->compression_cctx = NULL; + conn->compression_copy_started = false; + } + conn->compression_frame_ended = false; + if (conn->compression && !conn->compression_rejected && + strcmp(conn->compression, "off") != 0) + conn->compression_ready = true; + + /* No decoded row remains live at this boundary. */ + Assert(conn->decompressBuffer.start == conn->decompressBuffer.end); + conn->decompressBuffer.start = conn->decompressBuffer.cursor = + conn->decompressBuffer.end = 0; + if (conn->decompressBuffer.bufSize > PQ_COMPRESSION_BUFFER_KEEP_SIZE) + { + data = realloc(conn->decompressBuffer.buffer, 8192); + if (data) + { + conn->decompressBuffer.buffer = data; + conn->decompressBuffer.bufSize = 8192; + } + } + if (conn->compression_buffers_initialized && + conn->compression_output_buffer.maxlen > PQ_COMPRESSION_BUFFER_KEEP_SIZE) + { + data = realloc(conn->compression_output_buffer.data, INITIAL_EXPBUFFER_SIZE); + if (data) + { + conn->compression_output_buffer.data = data; + conn->compression_output_buffer.maxlen = INITIAL_EXPBUFFER_SIZE; + resetPQExpBuffer(&conn->compression_output_buffer); + } + } +#endif +} + +#ifdef USE_ZSTD +static int +pqCompressionInitBuffers(PGconn *conn) +{ + if (conn->compression_buffers_initialized) + return 0; + + initPQExpBuffer(&conn->compression_output_buffer); + if (PQExpBufferBroken(&conn->compression_output_buffer)) + { + termPQExpBuffer(&conn->compression_output_buffer); + libpq_append_conn_error(conn, "out of memory"); + return 1; + } + + conn->compression_buffers_initialized = true; + return 0; +} + +#endif + +/* + * Decode one complete outer message into the reusable msg_buffer. The caller + * consumes/traces the wrapper, then parses its ordinary messages separately. + */ +int +pqReadCompressedMessage(PGconn *conn, int msgLength) +{ +#ifdef USE_ZSTD + ZSTD_DCtx *dctx; + ZSTD_inBuffer input; + msg_buffer *buf = &conn->decompressBuffer; + size_t segment_size = 0; + size_t result = 1; + int position; + + if (!conn->compression_ready || conn->compression_frame_ended || + msgLength <= 0 || msgLength > PQ_COMPRESSION_MAX_WRAPPER_SIZE) + goto invalid; + + if (conn->compression_dctx == NULL) + { + size_t rc; + + dctx = ZSTD_createDCtx(); + if (dctx == NULL) + { + libpq_append_conn_error(conn, "out of memory"); + return -2; + } + rc = ZSTD_DCtx_setParameter(dctx, ZSTD_d_windowLogMax, PQ_COMPRESSION_WINDOW_LOG); + if (ZSTD_isError(rc)) + { + ZSTD_freeDCtx(dctx); + goto invalid; + } + conn->compression_dctx = dctx; + } + dctx = (ZSTD_DCtx *) conn->compression_dctx; + input.src = conn->inBuffer.buffer + conn->inBuffer.cursor; + input.size = msgLength; + input.pos = 0; + + for (;;) + { + ZSTD_outBuffer output; + size_t before = input.pos; + size_t room = PQ_COMPRESSION_MAX_SEGMENT_SIZE - segment_size; + + room = room == 0 ? 1 : Min(room, (size_t) 8192); + if (pqCheckMsgBufferSpace((size_t) buf->end + room, buf, conn)) + return -2; + output.dst = buf->buffer + buf->end; + output.size = room; + output.pos = 0; + result = ZSTD_decompressStream(dctx, &output, &input); + if (ZSTD_isError(result)) + goto invalid; + buf->end += output.pos; + segment_size += output.pos; + if (segment_size > PQ_COMPRESSION_MAX_SEGMENT_SIZE) + goto invalid; + if (result == 0) + { + if (input.pos != input.size) + goto invalid; + break; + } + if (input.pos == input.size && output.pos < output.size) + break; + if (input.pos == before && output.pos == 0) + goto invalid; + } + conn->compression_in_frame = (result != 0); + conn->compression_frame_ended = (result == 0); + + /* Validate the decoded prefix, including an unfinished message header. */ + position = buf->start; + while (buf->end - position >= 5) + { + uint32 len; + + if (buf->buffer[position] != PqMsg_DataRow && + buf->buffer[position] != PqMsg_CopyData) + goto invalid; + memcpy(&len, buf->buffer + position + 1, 4); + len = pg_ntoh32(len); + if (len < 4 || len >= INT_MAX) + goto invalid; + if (len > buf->end - position - 1) + break; + position += len + 1; + } + if (result == 0 && position != buf->end) + goto invalid; + conn->inBuffer.cursor += msgLength; + return 0; + +invalid: + libpq_append_conn_error(conn, "invalid compressed protocol message"); +#else + libpq_append_conn_error(conn, "client does not support compression with zstd"); +#endif + return -2; +} + +#ifdef USE_ZSTD +static int +pqPutCompressedCopySegment(PGconn *conn, const char *buffer, int nbytes, + ZSTD_EndDirective directive) +{ + ZSTD_CCtx *cctx = (ZSTD_CCtx *) conn->compression_cctx; + char header[5]; + bool context_advanced = false; + size_t source_size = nbytes > 0 ? (size_t) nbytes + 5 : 0; + size_t source_pos = 0; + size_t result = 0; + uint32 n32; + + if (cctx == NULL) + { + if (pqCompressionInitBuffers(conn)) + return EOF; + cctx = ZSTD_createCCtx(); + if (cctx == NULL) + goto oom; + result = ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, PQ_COMPRESSION_WINDOW_LOG); + if (ZSTD_isError(result)) + { + ZSTD_freeCCtx(cctx); + cctx = NULL; + goto zstd_error; + } + conn->compression_cctx = cctx; + } + + if (source_size > 0) + { + header[0] = PqMsg_CopyData; + n32 = pg_hton32((uint32) nbytes + 4); + memcpy(header + 1, &n32, 4); + } + do + { + ZSTD_EndDirective segment_directive; + ZSTD_outBuffer output; + size_t segment_size; + size_t segment_end; + size_t compressed_size; + + segment_size = Min(source_size - source_pos, + (size_t) PQ_COMPRESSION_MAX_SEGMENT_SIZE); + segment_end = source_pos + segment_size; + segment_directive = (segment_end == source_size) ? + directive : ZSTD_e_flush; + compressed_size = ZSTD_compressBound(segment_size); + /* Prefer sending pending data to enlarging the output buffer. */ + if ((size_t) conn->outCount + 5 + compressed_size > conn->outBufSize && + pqFlush(conn) < 0) + goto fail; + if (compressed_size > PQ_COMPRESSION_MAX_WRAPPER_SIZE || + compressed_size > INT_MAX - conn->outCount - 5 || + pqCheckOutBufferSpace(conn->outCount + 5 + compressed_size, conn)) + goto fail; + + resetPQExpBuffer(&conn->compression_output_buffer); + if (!enlargePQExpBuffer(&conn->compression_output_buffer, + compressed_size)) + goto oom; + output.dst = conn->compression_output_buffer.data; + output.size = compressed_size; + output.pos = 0; + + while (source_pos < segment_end) + { + ZSTD_inBuffer input; + size_t part_end; + + if (source_pos < sizeof(header)) + { + part_end = Min(segment_end, sizeof(header)); + input.src = header + source_pos; + input.size = part_end - source_pos; + } + else + { + part_end = segment_end; + input.src = buffer + source_pos - sizeof(header); + input.size = part_end - source_pos; + } + input.pos = 0; + while (input.pos < input.size) + { + context_advanced = true; + result = ZSTD_compressStream2(cctx, &output, &input, + ZSTD_e_continue); + if (ZSTD_isError(result)) + goto zstd_error; + if (output.pos == output.size && input.pos < input.size) + goto output_too_small; + } + source_pos = part_end; + } + + { + ZSTD_inBuffer input = {NULL, 0, 0}; + + do + { + context_advanced = true; + result = ZSTD_compressStream2(cctx, &output, &input, + segment_directive); + if (ZSTD_isError(result)) + goto zstd_error; + if (output.pos == output.size && result != 0) + goto output_too_small; + } while (result != 0); + } + + if (pqPutMsgStart(PqMsg_CompressedData, conn) < 0 || + pqPutnchar(conn->compression_output_buffer.data, output.pos, conn) < 0 || + pqPutMsgEnd(conn) < 0) + goto fail; + } while (source_pos < source_size); + + return 0; + +oom: + libpq_append_conn_error(conn, "out of memory"); + goto fail; + +zstd_error: + libpq_append_conn_error(conn, "Zstandard compression failed: %s", + ZSTD_getErrorName(result)); + goto fail; + +output_too_small: + libpq_append_conn_error(conn, "Zstandard compression output buffer is too small"); +fail: + if (context_advanced) + conn->status = CONNECTION_BAD; + return EOF; +} + +int +pqPutCompressedCopyData(PGconn *conn, const char *buffer, int nbytes) +{ + if (pqPutCompressedCopySegment(conn, buffer, nbytes, ZSTD_e_flush) < 0) + return EOF; + conn->compression_copy_started = true; + return 0; +} + +int +pqEndCompressedCopyData(PGconn *conn) +{ + if (!conn->compression_copy_started) + return 0; + if (pqPutCompressedCopySegment(conn, NULL, 0, ZSTD_e_end) < 0) + return EOF; + conn->compression_copy_started = false; + return 0; +} + +#endif diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c index 0941ae4b6ca..f4118536efe 100644 --- a/src/interfaces/libpq/fe-connect.c +++ b/src/interfaces/libpq/fe-connect.c @@ -416,6 +416,10 @@ static const internalPQconninfoOption PQconninfoOptions[] = { "SSL-Key-Log-File", "D", 64, offsetof(struct pg_conn, sslkeylogfile)}, + {"compression", "PGCOMPRESSION", "off", NULL, + "Protocol-Compression", "", 6, + offsetof(struct pg_conn, compression)}, + /* Terminating entry --- MUST BE LAST */ {NULL, NULL, NULL, NULL, NULL, NULL, 0} @@ -541,7 +545,10 @@ pqDropConnection(PGconn *conn, bool flushInput) /* Optionally discard any unread data */ if (flushInput) + { conn->inBuffer.start = conn->inBuffer.cursor = conn->inBuffer.end = 0; + pqCompressionReset(conn); + } /* Always discard any unsent data */ conn->outCount = 0; @@ -2175,6 +2182,23 @@ pqConnectOptions2(PGconn *conn) goto oom_error; } + if (conn->compression && + strcmp(conn->compression, "off") != 0 && + strcmp(conn->compression, "prefer") != 0 && + strcmp(conn->compression, "zstd") != 0) + { + libpq_append_conn_error(conn, "invalid compression value: \"%s\"", + conn->compression); + return false; + } +#ifndef USE_ZSTD + if (conn->compression && strcmp(conn->compression, "zstd") == 0) + { + libpq_append_conn_error(conn, "compression value \"zstd\" is not supported"); + return false; + } +#endif + /* * Only if we get this far is it appropriate to try to connect. (We need a * state flag, rather than just the boolean result of this function, in @@ -2739,7 +2763,8 @@ pqConnectDBStart(PGconn *conn) goto connect_errReturn; } - /* Ensure our buffers are empty */ + /* Ensure buffers and compression state are empty on every connection. */ + pqCompressionReset(conn); conn->inBuffer.start = conn->inBuffer.cursor = conn->inBuffer.end = 0; conn->outCount = 0; @@ -5131,6 +5156,7 @@ freePGconn(PGconn *conn) free(conn->scram_client_key); free(conn->scram_server_key); free(conn->sslkeylogfile); + free(conn->compression); free(conn->oauth_issuer); free(conn->oauth_issuer_id); free(conn->oauth_discovery_uri); @@ -5150,6 +5176,7 @@ freePGconn(PGconn *conn) free(conn->inBuffer.buffer); free(conn->outBuffer); free(conn->rowBuf); + pqCompressionReset(conn); termPQExpBuffer(&conn->errorMessage); termPQExpBuffer(&conn->workBuffer); diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c index 294690648bd..4c743d7fdc4 100644 --- a/src/interfaces/libpq/fe-exec.c +++ b/src/interfaces/libpq/fe-exec.c @@ -29,6 +29,11 @@ #include "libpq-int.h" #include "mb/pg_wchar.h" +#ifdef USE_ZSTD +/* Avoid per-message compression and flush overhead for very small inputs. */ +#define PQ_COMPRESSION_MIN_INPUT_SIZE 1024 +#endif + /* keep this in same order as ExecStatusType in libpq-fe.h */ char *const pgresStatus[] = { "PGRES_EMPTY_QUERY", @@ -2732,6 +2737,17 @@ PQputCopyData(PGconn *conn, const char *buffer, int nbytes) if (nbytes > 0) { +#ifdef USE_ZSTD + if (conn->compression_ready && + conn->asyncStatus == PGASYNC_COPY_IN && + nbytes >= PQ_COMPRESSION_MIN_INPUT_SIZE) + { + if (pqPutCompressedCopyData(conn, buffer, nbytes) < 0) + return -1; + return 1; + } +#endif + /* * Try to flush any previously sent data in preference to growing the * output buffer. If we can't enlarge the buffer enough to hold the @@ -2774,6 +2790,13 @@ PQputCopyEnd(PGconn *conn, const char *errormsg) return -1; } +#ifdef USE_ZSTD + if (conn->compression_ready && + conn->asyncStatus == PGASYNC_COPY_IN && + pqEndCompressedCopyData(conn) < 0) + return -1; +#endif + /* * Send the COPY END indicator. This is simple enough that we don't * bother delegating it to the fe-protocol files. diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c index 2bb8b6e67f4..ec7792c1bd6 100644 --- a/src/interfaces/libpq/fe-misc.c +++ b/src/interfaces/libpq/fe-misc.c @@ -446,7 +446,8 @@ pqParseDone(PGconn *conn, msg_buffer *msgBuf, int newInStart) { /* trace server-to-client message */ if (conn->Pfdebug) - pqTraceOutputMessage(conn, msgBuf->buffer + msgBuf->start, false); + pqTraceOutputMessage(conn, msgBuf->buffer + msgBuf->start, false, + msgBuf == &conn->decompressBuffer); /* Mark message as done */ msgBuf->start = newInStart; @@ -546,7 +547,7 @@ pqPutMsgEnd(PGconn *conn) if (conn->Pfdebug) { if (conn->outCount < conn->outMsgStart) - pqTraceOutputMessage(conn, conn->outBuffer + conn->outCount, true); + pqTraceOutputMessage(conn, conn->outBuffer + conn->outCount, true, false); else pqTraceOutputNoTypeByteMessage(conn, conn->outBuffer + conn->outMsgStart); diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c index 63a308915c9..dafd7755f6b 100644 --- a/src/interfaces/libpq/fe-protocol3.c +++ b/src/interfaces/libpq/fe-protocol3.c @@ -43,6 +43,7 @@ (id) == PqMsg_NoticeResponse || \ (id) == PqMsg_NotificationResponse || \ (id) == PqMsg_RowDescription || \ + (id) == PqMsg_CompressedData || \ (id) == PqMsg_ParameterDescription) @@ -61,6 +62,16 @@ static void reportErrorPosition(PQExpBuffer msg, const char *query, static size_t build_startup_packet(const PGconn *conn, char *packet, const PQEnvironmentOption *options); +/* Select the input buffer and report invalid compressed input. */ +static msg_buffer * +getMsgBuffer(PGconn *conn) +{ + msg_buffer *buf = pqGetMessageBuffer(conn); + + if (buf == NULL) + handleFatalError(conn); + return buf; +} /* * parseInput: if appropriate, parse input data from backend @@ -73,13 +84,17 @@ pqParseInput3(PGconn *conn) char id; int msgLength; int avail; - msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *msgBuf; /* * Loop to parse successive complete messages available in the buffer. */ for (;;) { + msgBuf = getMsgBuffer(conn); + if (!msgBuf) + return; + /* * Try to read a message. First get the type code and length. Return * if not enough data. @@ -285,6 +300,13 @@ pqParseInput3(PGconn *conn) case PqMsg_BindComplete: /* Nothing to do for this message type */ break; + case PqMsg_CompressedData: + if (pqReadCompressedMessage(conn, msgLength) != 0) + { + handleFatalError(conn); + return; + } + break; case PqMsg_CloseComplete: /* If we're doing PQsendClose, we're done; else ignore */ if (conn->cmd_queue_head && @@ -1512,16 +1534,25 @@ pqGetNegotiateProtocolVersion3(PGconn *conn, msg_buffer *msgBuf) /* the version is acceptable */ conn->pversion = their_version; - /* - * We don't currently request any protocol extensions, so we don't expect - * the server to reply with any either. - */ + /* Recognize rejection of our optional compression profile. */ for (int i = 0; i < num; i++) { if (pqGets(&conn->workBuffer, conn, msgBuf)) { goto eof; } +#ifdef USE_ZSTD + if (strcmp(conn->workBuffer.data, "_pq_.compression") == 0 && + conn->compression && strcmp(conn->compression, "off") != 0 && + !conn->compression_rejected) + { + conn->compression_rejected = true; + if (strcmp(conn->compression, "prefer") == 0) + continue; + libpq_append_conn_error(conn, "server does not support protocol compression method \"zstd\""); + goto failure; + } +#endif if (strncmp(conn->workBuffer.data, "_pq_.", 5) != 0) { libpq_append_conn_error(conn, "received invalid protocol negotiation message: server reported unsupported parameter name without a \"%s\" prefix (\"%s\")", "_pq_.", conn->workBuffer.data); @@ -1796,6 +1827,7 @@ getReadyForQuery(PGconn *conn, msg_buffer *msgBuf) break; } + pqCompressionReady(conn); return 0; } @@ -1811,12 +1843,15 @@ getCopyDataMessage(PGconn *conn, msg_buffer **out_buf) char id; int msgLength; int avail; - msg_buffer *msgBuf = &conn->inBuffer; - - *out_buf = msgBuf; + msg_buffer *msgBuf; for (;;) { + msgBuf = getMsgBuffer(conn); + if (!msgBuf) + return -2; + *out_buf = msgBuf; + /* * Do we have the next input message? To make life simpler for async * callers, we keep returning 0 until the next message is fully @@ -1866,6 +1901,13 @@ getCopyDataMessage(PGconn *conn, msg_buffer **out_buf) if (getNotify(conn, msgBuf)) return 0; break; + case PqMsg_CompressedData: + if (pqReadCompressedMessage(conn, msgLength - 4) != 0) + { + handleFatalError(conn); + return -2; + } + break; case PqMsg_NoticeResponse: if (pqGetErrorNotice3(conn, false, msgBuf)) return 0; @@ -1924,7 +1966,7 @@ int pqGetCopyData3(PGconn *conn, char **buffer, int async) { int msgLength; - msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *msgBuf; for (;;) { @@ -2036,7 +2078,7 @@ pqGetlineAsync3(PGconn *conn, char *buffer, int bufsize) { int msgLength; int avail; - msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *msgBuf; if (conn->asyncStatus != PGASYNC_COPY_OUT && conn->asyncStatus != PGASYNC_COPY_BOTH) @@ -2100,6 +2142,13 @@ pqEndcopy3(PGconn *conn) return 1; } +#ifdef USE_ZSTD + if (conn->compression_ready && + conn->asyncStatus == PGASYNC_COPY_IN && + pqEndCompressedCopyData(conn) < 0) + return 1; +#endif + /* Send the CopyDone message if needed */ if (conn->asyncStatus == PGASYNC_COPY_IN || conn->asyncStatus == PGASYNC_COPY_BOTH) @@ -2192,7 +2241,7 @@ pqFunctionCall3(PGconn *conn, Oid fnid, int msgLength; int avail; int i; - msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *msgBuf; /* already validated by PQnfn */ Assert(conn->pipelineStatus == PQ_PIPELINE_OFF); @@ -2250,6 +2299,9 @@ pqFunctionCall3(PGconn *conn, Oid fnid, */ needInput = true; + msgBuf = getMsgBuffer(conn); + if (!msgBuf) + break; msgBuf->cursor = msgBuf->start; if (pqGetc(&id, conn, msgBuf)) continue; @@ -2342,6 +2394,13 @@ pqFunctionCall3(PGconn *conn, Oid fnid, continue; status = PGRES_FATAL_ERROR; break; + case PqMsg_CompressedData: + if (pqReadCompressedMessage(conn, msgLength) != 0) + { + handleFatalError(conn); + return pqPrepareAsyncResult(conn); + } + break; case PqMsg_NotificationResponse: /* handle notify and go back to processing return values */ if (getNotify(conn, msgBuf)) @@ -2510,6 +2569,11 @@ build_startup_packet(const PGconn *conn, char *packet, if (conn->client_encoding_initial && conn->client_encoding_initial[0]) ADD_STARTUP_OPTION("client_encoding", conn->client_encoding_initial); +#ifdef USE_ZSTD + if (conn->compression && strcmp(conn->compression, "off") != 0) + ADD_STARTUP_OPTION("_pq_.compression", "zstd"); +#endif + /* Add any environment-driven GUC settings needed */ for (next_eo = options; next_eo->envName; next_eo++) { diff --git a/src/interfaces/libpq/fe-trace.c b/src/interfaces/libpq/fe-trace.c index 2901fa5b451..dcfd6a327e1 100644 --- a/src/interfaces/libpq/fe-trace.c +++ b/src/interfaces/libpq/fe-trace.c @@ -622,11 +622,12 @@ pqTraceOutput_ReadyForQuery(FILE *f, const char *message, int *cursor) * Print the given message to the trace output stream. */ void -pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer) +pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer, + bool compressed) { char id; int length; - char *prefix = toServer ? "F" : "B"; + const char *prefix = toServer ? "F" : compressed ? "B*" : "B"; int logCursor = 0; bool regress; @@ -651,7 +652,8 @@ pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer) * name) fields can change as server code is modified, and if their * lengths differ from the originals, that would break tests. */ - if (regress && !toServer && (id == PqMsg_ErrorResponse || id == PqMsg_NoticeResponse)) + if (regress && (id == PqMsg_CompressedData || + (!toServer && (id == PqMsg_ErrorResponse || id == PqMsg_NoticeResponse)))) fprintf(conn->Pfdebug, "%s\tNN\t", prefix); else fprintf(conn->Pfdebug, "%s\t%d\t", prefix, length); @@ -688,6 +690,11 @@ pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer) else pqTraceOutput_CommandComplete(conn->Pfdebug, message, &logCursor); break; + case PqMsg_CompressedData: + fprintf(conn->Pfdebug, "CompressedData"); + /* Trace decoded messages separately instead of the payload. */ + logCursor = length + 1; + break; case PqMsg_CopyData: pqTraceOutput_CopyData(conn->Pfdebug, message, &logCursor, length, regress); diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h index 673648fa154..2164f111deb 100644 --- a/src/interfaces/libpq/libpq-int.h +++ b/src/interfaces/libpq/libpq-int.h @@ -375,6 +375,7 @@ typedef struct msg_buffer int end; /* offset to first position after avail data */ } msg_buffer; + /* * PGconn stores all the state data associated with a single connection * to a backend. @@ -439,6 +440,7 @@ struct pg_conn char *scram_client_key; /* base64-encoded SCRAM client key */ char *scram_server_key; /* base64-encoded SCRAM server key */ char *sslkeylogfile; /* where should the client write ssl keylogs */ + char *compression; /* protocol compression policy */ bool cancelRequest; /* true if this connection is used to send a * cancel request, instead of being a normal @@ -582,7 +584,16 @@ struct pg_conn * code that uses ints during size calculations. */ msg_buffer inBuffer; - + msg_buffer decompressBuffer; /* decompressed ordinary messages */ + void *compression_dctx; + void *compression_cctx; + PQExpBufferData compression_output_buffer; + bool compression_buffers_initialized; + bool compression_in_frame; + bool compression_frame_ended; + bool compression_ready; + bool compression_copy_started; + bool compression_rejected; /* Buffer for data not yet sent to backend */ char *outBuffer; /* currently allocated buffer */ @@ -793,6 +804,16 @@ extern PGresult *pqFunctionCall3(PGconn *conn, Oid fnid, extern int PQsendCancelRequest(PGconn *cancelConn); +/* === in fe-compress.c === */ +extern msg_buffer *pqGetMessageBuffer(PGconn *conn); +extern int pqReadCompressedMessage(PGconn *conn, int msgLength); +extern void pqCompressionReset(PGconn *conn); +extern void pqCompressionReady(PGconn *conn); +#ifdef USE_ZSTD +extern int pqPutCompressedCopyData(PGconn *conn, const char *buffer, int nbytes); +extern int pqEndCompressedCopyData(PGconn *conn); +#endif + /* === in fe-misc.c === */ /* @@ -921,7 +942,7 @@ extern ssize_t pg_GSS_bytes_pending(PGconn *conn); /* === in fe-trace.c === */ extern void pqTraceOutputMessage(PGconn *conn, const char *message, - bool toServer); + bool toServer, bool compressed); extern void pqTraceOutputNoTypeByteMessage(PGconn *conn, const char *message); extern void pqTraceOutputCharResponse(PGconn *conn, const char *responseType, char response); diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build index b0ae72167a1..9e2f66b7227 100644 --- a/src/interfaces/libpq/meson.build +++ b/src/interfaces/libpq/meson.build @@ -5,6 +5,7 @@ libpq_sources = files( 'fe-auth-scram.c', 'fe-auth.c', 'fe-cancel.c', + 'fe-compress.c', 'fe-connect.c', 'fe-exec.c', 'fe-lobj.c', @@ -161,6 +162,7 @@ tests += { 't/004_load_balance_dns.pl', 't/005_negotiate_encryption.pl', 't/006_service.pl', + 't/007_compression.pl', ], 'env': { 'with_ssl': ssl_library, diff --git a/src/interfaces/libpq/nls.mk b/src/interfaces/libpq/nls.mk index 3fa87a0aaac..692c10e0939 100644 --- a/src/interfaces/libpq/nls.mk +++ b/src/interfaces/libpq/nls.mk @@ -4,6 +4,7 @@ GETTEXT_FILES = fe-auth.c \ fe-auth-oauth.c \ fe-auth-scram.c \ fe-cancel.c \ + fe-compress.c \ fe-connect.c \ fe-exec.c \ fe-gssapi-common.c \ diff --git a/src/interfaces/libpq/t/007_compression.pl b/src/interfaces/libpq/t/007_compression.pl new file mode 100644 index 00000000000..ccded93fda2 --- /dev/null +++ b/src/interfaces/libpq/t/007_compression.pl @@ -0,0 +1,232 @@ +# Copyright (c) 2026, PostgreSQL Global Development Group + +use strict; +use warnings FATAL => 'all'; + +use IPC::Run; +use PostgreSQL::Test::Cluster; +use PostgreSQL::Test::Utils; +use Test::More; + +plan skip_all => 'Zstandard is not supported by this build' + unless check_pg_config('#define USE_ZSTD 1'); + +# Comparing the output alone would miss a failure shared by both clients. +sub run_checked +{ + my ($command, $expected_status) = @_; + $expected_status //= 0; + my ($stdout, $stderr); + IPC::Run::run($command, '>', \$stdout, '2>', \$stderr); + my $status = $?; + is($status, $expected_status << 8, 'psql exits with the expected status') + or diag($stderr); + chomp($stdout, $stderr); + return ($stdout, $stderr); +} + +my $node = PostgreSQL::Test::Cluster->new('node'); +$node->init(allows_streaming => 1); +$node->append_conf('postgresql.conf', 'protocol_compression = zstd'); +$node->append_conf('postgresql.conf', 'summarize_wal = on'); +$node->start; + +my @commands = ( + '-c', 'SELECT 1', + '-c', q{SELECT NULL::text, ''::text, E'a\nb', true, -1::bigint, + 1.25::numeric, '\x00017fff'::bytea}, + '-c', q{SELECT g, CASE WHEN g % 3 = 0 THEN NULL + ELSE repeat(chr(64 + g), g) END FROM generate_series(1, 26) AS g}, + '-c', 'SELECT g FROM generate_series(1, 2000) AS g', + '-c', q{SELECT repeat('x', 1024 * 1024)}, + '-c', q{DO $$ BEGIN RAISE NOTICE 'compression notice'; END $$}, + '-c', 'BEGIN', + '-c', q{SELECT 'in transaction', repeat('y', 8192)}, + '-c', 'COMMIT', + '-c', q{COPY (SELECT repeat('copy data ', 20) FROM generate_series(1, 100)) TO STDOUT}, + '-c', 'SELECT 2'); + +my ($plain_stdout, $plain_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=off', @commands ]); +my ($compressed_stdout, $compressed_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', @commands ]); +is($compressed_stdout, $plain_stdout, + 'compressed queries and COPY OUT produce the same output'); +is($compressed_stderr, $plain_stderr, + 'compressed and uncompressed connections report the same notices'); + +my @error_commands = ( + '-c', q{SELECT 'before error', repeat('a', 8192)}, + '-c', 'SELECT 1 / 0', + '-c', q{SELECT 'after error', repeat('b', 8192)}); +my ($plain_error_stdout, $plain_error_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=off', + @error_commands ]); +my ($compressed_error_stdout, $compressed_error_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', + @error_commands ]); +is($compressed_error_stdout, $plain_error_stdout, + 'compression preserves results around an error response'); +is($compressed_error_stderr, $plain_error_stderr, + 'compression preserves an error response'); + +my $copy_file = $node->basedir . '/copy.data'; +append_to_file($copy_file, + join('', map { "$_\tcopy data $_\n" } 1 .. 2000)); +my ($copy_in_stdout, $copy_in_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', + '-c', 'CREATE TABLE copy_in_test (id integer, value text)', + '-c', "\\copy copy_in_test FROM '$copy_file'", + '-c', 'SELECT count(*), min(id), max(id) FROM copy_in_test' ]); +is($copy_in_stderr, '', 'compressed COPY IN produced no errors'); +like($copy_in_stdout, qr/2000\|1\|2000/, + 'compressed COPY IN loaded all rows'); + +my $bad_copy_file = $node->basedir . '/copy-error.data'; +append_to_file($bad_copy_file, + "invalid integer\tfirst row\n" . slurp_file($copy_file)); +my ($copy_error_stdout, $copy_error_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', + '-c', "\\copy copy_in_test FROM '$bad_copy_file'", + '-c', 'SELECT 42' ]); +like($copy_error_stderr, qr/invalid input syntax for type integer/, + 'compressed COPY reports a row error'); +is($copy_error_stdout, '42', + 'connection remains usable after discarding failed COPY input'); + +my $binary_file = $node->basedir . '/copy.binary'; +$node->command_ok( + [ 'psql', '-XAt', '--dbname', + $node->connstr('postgres') . ' compression=zstd', + '-c', "\\copy (SELECT g, md5(g::text) FROM generate_series(1, 2000) g) TO '$binary_file' (FORMAT binary)" ], + 'compressed binary COPY OUT'); +$node->safe_psql('postgres', + 'CREATE TABLE binary_copy_test (id integer, value text)'); +$node->command_ok( + [ 'psql', '-XAt', '--dbname', + $node->connstr('postgres') . ' compression=zstd', + '-c', "\\copy binary_copy_test FROM '$binary_file' (FORMAT binary)" ], + 'compressed binary COPY IN'); +is($node->safe_psql('postgres', + 'SELECT count(*), min(id), max(id), bool_and(value = md5(id::text)) FROM binary_copy_test'), + '2000|1|2000|t', 'compressed binary COPY preserved all values'); + +my $pgbench_script = $node->basedir . '/compression.pgbench'; +append_to_file($pgbench_script, <<'EOS'); +SELECT 42 AS answer, repeat('extended result ', 5000) AS payload \gset +\if :answer != 42 + SELECT 1 / 0; +\endif +EOS +for my $query_mode ('extended', 'prepared') +{ + $node->command_ok( + [ 'pgbench', '--no-vacuum', '--client=1', '--transactions=3', + '--protocol', $query_mode, '--file', $pgbench_script, + $node->connstr('postgres') . ' compression=zstd' ], + "compressed $query_mode query protocol"); +} + +$node->command_ok( + [ 'libpq_testclient', '--compression-reset', + $node->connstr('postgres') . ' compression=zstd' ], + 'PQreset discards partial compression state and starts a new stream'); + +$node->command_ok( + [ 'libpq_testclient', '--compression-large-datarow', + $node->connstr('postgres') . ' compression=zstd' ], + 'large backend DataRow spans compressed segments'); + +$node->command_ok( + [ 'libpq_testclient', '--compression-large-copy', + $node->connstr('postgres') . ' compression=zstd' ], + 'large frontend CopyData, tracing, and reuse with mixed COPY messages'); + +$node->command_ok( + [ 'libpq_testclient', '--compression-incompressible-copy', + $node->connstr('postgres') . ' compression=zstd' ], + 'incompressible COPY data crosses the segment limit in both directions'); + +$node->command_ok( + [ 'libpq_testclient', '--compression-frame-boundaries', + $node->connstr('postgres') . + ' compression=zstd sslmode=disable gssencmode=disable' ], + 'frontend compression enforces Zstandard frame boundaries'); + +$node->command_ok( + [ 'libpq_testclient', '--compression-pipeline', + $node->connstr('postgres') . ' compression=zstd' ], + 'compressed server responses work in libpq pipeline mode'); + +my $trace_file = $node->basedir . '/compression.trace'; +$node->command_ok( + [ 'libpq_testclient', '--compression-trace', + $node->connstr('postgres') . ' compression=zstd', $trace_file ], + 'compressed frames end at ReadyForQuery, including inside a transaction'); +my $trace = slurp_file($trace_file); +like($trace, qr/^B\t\d+\tCompressedData$/m, + 'ordinary trace omits the compressed payload'); +like($trace, qr/^B\tNN\tCompressedData/m, 'trace shows compressed wrappers'); +like($trace, qr/^B\*\t\d+\tDataRow/m, 'trace shows decoded rows'); +like($trace, qr/^B\*\t\d+\tCopyData/m, 'trace shows decoded COPY messages'); +unlike($trace, qr/^B\*\t\S+\t(?!DataRow|CopyData)/m, + 'control messages stay outside compression'); +unlike($trace, qr/mismatched message length/, 'tracing accounts for each message'); + +my $dump_file = $node->basedir . '/copy_in_test.dump'; +$node->command_ok( + [ 'pg_dump', '--format=custom', '--file', $dump_file, + '--table=copy_in_test', $node->connstr('postgres') ], + 'created archive for compressed pg_restore'); +$node->safe_psql('postgres', 'DROP TABLE copy_in_test'); +$node->command_ok( + [ 'pg_restore', '--dbname', + $node->connstr('postgres') . ' compression=zstd', $dump_file ], + 'pg_restore uses compressed COPY IN'); +is($node->safe_psql('postgres', + 'SELECT count(*), min(id), max(id) FROM copy_in_test'), + '2000|1|2000', 'compressed pg_restore loaded all rows'); + +# An incremental backup sends its manifest with COPY, which is the only +# frontend COPY the walsender reads, and it ends the compressed frame with an +# empty CompressedData message. +my $full_backup = $node->basedir . '/full_backup'; +$node->command_ok( + [ 'pg_basebackup', '--no-sync', '--checkpoint' => 'fast', + '--pgdata' => $full_backup, + '--dbname' => $node->connstr('postgres') ], + 'full backup for a compressed UPLOAD_MANIFEST'); +$node->safe_psql('postgres', + 'CREATE TABLE after_backup AS SELECT g FROM generate_series(1, 1000) g'); +my $incremental_backup = $node->basedir . '/incremental_backup'; +$node->command_ok( + [ 'pg_basebackup', '--no-sync', '--checkpoint' => 'fast', + '--pgdata' => $incremental_backup, + '--incremental' => $full_backup . '/backup_manifest', + '--dbname' => $node->connstr('postgres') . ' compression=zstd' ], + 'incremental backup uploads its manifest over a compressed connection'); + +my (undef, $invalid_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=invalid', + '-c', 'SELECT 1' ], 2); +like($invalid_stderr, qr/invalid compression value: "invalid"/, + 'invalid compression value is rejected'); + +$node->safe_psql('postgres', "ALTER SYSTEM SET protocol_compression = 'off'"); +$node->reload; +ok($node->poll_query_until('postgres', + q{SELECT current_setting('protocol_compression') = 'off'}), + 'protocol compression is disabled after reload'); +my ($preferred_stdout, $preferred_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=prefer', + '-c', 'SELECT 1' ]); +is($preferred_stderr, '', 'preferred compression falls back without errors'); +is($preferred_stdout, '1', 'preferred compression fallback returns output'); +my (undef, $disabled_stderr) = run_checked( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', + '-c', 'SELECT 1' ], 2); +like($disabled_stderr, qr/does not support protocol compression method "zstd"/, + 'server can reject protocol compression'); + +$node->stop('fast'); +done_testing(); diff --git a/src/interfaces/libpq/test/libpq_testclient.c b/src/interfaces/libpq/test/libpq_testclient.c index 20730709ee7..a3cd4f94ca7 100644 --- a/src/interfaces/libpq/test/libpq_testclient.c +++ b/src/interfaces/libpq/test/libpq_testclient.c @@ -1,6 +1,6 @@ /* * libpq_testclient.c - * A test program for the libpq public API + * A test program for libpq and the frontend/backend protocol * * Copyright (c) 2022-2026, PostgreSQL Global Development Group * @@ -10,8 +10,17 @@ #include "postgres_fe.h" +#include + #include "libpq-fe.h" +#ifdef USE_ZSTD +#include "libpq/protocol_compression.h" + +#include "libpq-int.h" +#include "port/pg_bswap.h" +#endif + static void print_ssl_library(void) { @@ -23,6 +32,674 @@ print_ssl_library(void) printf("%s\n", lib); } +static int +test_compression_reset(const char *conninfo) +{ + PGconn *conn = PQconnectdb(conninfo); + PGresult *res; + + if (PQstatus(conn) != CONNECTION_OK) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return 1; + } +#ifdef USE_ZSTD + res = PQexec(conn, "SELECT repeat('before reset ', 1000)"); + if (PQresultStatus(res) != PGRES_TUPLES_OK) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQclear(res); + PQfinish(conn); + return 1; + } + PQclear(res); + if (conn->compression_dctx == NULL || + conn->decompressBuffer.start != conn->decompressBuffer.end) + { + fprintf(stderr, "server response was not compressed\n"); + PQfinish(conn); + return 1; + } + + /* + * Model an unfinished compressed DataRow without relying on socket + * timing. + */ + if (conn->decompressBuffer.buffer == NULL || + conn->decompressBuffer.bufSize < 1) + { + fprintf(stderr, "out of memory\n"); + PQfinish(conn); + return 1; + } + conn->decompressBuffer.buffer[0] = PqMsg_DataRow; + conn->decompressBuffer.start = 0; + conn->decompressBuffer.end = 1; + conn->compression_in_frame = true; +#else + res = PQexec(conn, "SELECT repeat('x', 17 * 1024 * 1024)"); + if (PQresultStatus(res) != PGRES_TUPLES_OK || + PQgetlength(res, 0, 0) != 17 * 1024 * 1024) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQclear(res); + PQfinish(conn); + return 1; + } + PQclear(res); +#endif + PQreset(conn); + if (PQstatus(conn) != CONNECTION_OK) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return 1; + } +#ifdef USE_ZSTD + if (conn->compression_dctx != NULL || + conn->compression_buffers_initialized || + conn->compression_in_frame) + { + fprintf(stderr, "PQreset did not discard compression state\n"); + PQfinish(conn); + return 1; + } +#endif + res = PQexec(conn, "SELECT repeat('after reset ', 1000)"); + if (PQresultStatus(res) != PGRES_TUPLES_OK) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQclear(res); + PQfinish(conn); + return 1; + } + PQclear(res); +#ifdef USE_ZSTD + if (conn->compression_dctx == NULL) + { + fprintf(stderr, "server response after PQreset was not compressed\n"); + PQfinish(conn); + return 1; + } +#endif + PQfinish(conn); + return 0; +} + +#ifdef USE_ZSTD +static PGconn * +compression_connect(const char *conninfo) +{ + PGconn *conn = PQconnectdb(conninfo); + + if (PQstatus(conn) != CONNECTION_OK) + { + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return NULL; + } + if (PQsslInUse(conn) || PQgssEncInUse(conn)) + { + fprintf(stderr, "raw protocol test requires an unencrypted connection\n"); + PQfinish(conn); + return NULL; + } + return conn; +} + +static int +test_compression_large_datarow(const char *conninfo) +{ + const int data_size = 17 * 1024 * 1024; + PGconn *conn; + PGresult *res = NULL; + const char *data; + int i; + int status = 1; + + conn = compression_connect(conninfo); + if (conn == NULL) + return 1; + + res = PQexec(conn, "SELECT repeat('x', 17 * 1024 * 1024)"); + if (PQresultStatus(res) != PGRES_TUPLES_OK || + PQntuples(res) != 1 || PQnfields(res) != 1 || + PQgetisnull(res, 0, 0) || PQgetlength(res, 0, 0) != data_size) + goto done; + + data = PQgetvalue(res, 0, 0); + for (i = 0; i < data_size; i++) + { + if (data[i] != 'x') + goto done; + } + + /* Creating the decompressor proves that a wrapper was received. */ + if (conn->compression_dctx == NULL) + goto done; + + status = 0; + +done: + if (status != 0) + fprintf(stderr, "large compressed DataRow was not preserved: %s", + PQerrorMessage(conn)); + PQclear(res); + PQfinish(conn); + return status; +} + +static int +send_raw_bytes(PGconn *conn, const char *data, size_t len) +{ + int sock = PQsocket(conn); + + while (len > 0) + { + int flags = 0; + int sent; + +#ifdef MSG_NOSIGNAL + flags |= MSG_NOSIGNAL; +#endif + sent = send(sock, data, (int) Min(len, (size_t) INT_MAX), flags); + if (sent <= 0) + return 1; + data += sent; + len -= sent; + } + return 0; +} + +/* Send a complete or deliberately unfinished frame containing one CopyData. */ +static int +send_compressed_copy_frame(PGconn *conn, const char *data, size_t len, + bool finish) +{ + char *input_data; + char *wire_data; + size_t input_size = len + 5; + size_t output_size = ZSTD_compressBound(input_size) + + ZSTD_CStreamOutSize() + 64; + ZSTD_CCtx *cctx = NULL; + ZSTD_inBuffer input; + ZSTD_outBuffer output; + size_t result; + uint32 n32; + int status = 1; + + input_data = malloc(input_size); + wire_data = malloc(output_size + 5); + if (input_data == NULL || wire_data == NULL) + goto done; + input_data[0] = PqMsg_CopyData; + n32 = pg_hton32((uint32) len + 4); + memcpy(input_data + 1, &n32, 4); + memcpy(input_data + 5, data, len); + + cctx = ZSTD_createCCtx(); + if (cctx == NULL) + goto done; + result = ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, 16); + if (ZSTD_isError(result)) + goto done; + + input.src = input_data; + input.size = input_size; + input.pos = 0; + output.dst = wire_data + 5; + output.size = output_size; + output.pos = 0; + do + { + result = ZSTD_compressStream2(cctx, &output, &input, + finish ? ZSTD_e_end : ZSTD_e_flush); + if (ZSTD_isError(result)) + goto done; + } while (result != 0); + + wire_data[0] = PqMsg_CompressedData; + n32 = pg_hton32((uint32) output.pos + 4); + memcpy(wire_data + 1, &n32, 4); + status = send_raw_bytes(conn, wire_data, output.pos + 5); + +done: + ZSTD_freeCCtx(cctx); + free(input_data); + free(wire_data); + return status; +} + +static int +start_copy(PGconn *conn) +{ + PGresult *res; + + res = PQexec(conn, "CREATE TEMP TABLE compression_frame (data text); " + "COPY compression_frame FROM STDIN"); + if (PQresultStatus(res) != PGRES_COPY_IN) + { + PQclear(res); + return 1; + } + PQclear(res); + return 0; +} + +static int +expect_copy_result(PGconn *conn, ExecStatusType expected) +{ + PGresult *res; + bool matched = false; + bool unexpected = false; + + while ((res = PQgetResult(conn)) != NULL) + { + if (PQresultStatus(res) == expected) + matched = true; + else + unexpected = true; + PQclear(res); + } + return matched && !unexpected ? 0 : 1; +} + +static int +expect_protocol_rejection(PGconn *conn) +{ + if (PQputCopyEnd(conn, NULL) != 1) + return PQstatus(conn) == CONNECTION_BAD ? 0 : 1; + return expect_copy_result(conn, PGRES_FATAL_ERROR); +} + +static int +test_compression_frame_boundaries(const char *conninfo) +{ + PGconn *conn; + PGresult *res; + int status = 1; + + /* The frame epilogue may share a wrapper with the final CopyData. */ + conn = compression_connect(conninfo); + if (conn == NULL) + return 1; + if (start_copy(conn) || + send_compressed_copy_frame(conn, "complete\n", strlen("complete\n"), + true) || + PQputCopyEnd(conn, NULL) != 1 || + expect_copy_result(conn, PGRES_COMMAND_OK)) + goto done; + res = PQexec(conn, "SELECT data FROM compression_frame"); + if (PQresultStatus(res) != PGRES_TUPLES_OK || + PQntuples(res) != 1 || strcmp(PQgetvalue(res, 0, 0), "complete") != 0) + { + PQclear(res); + goto done; + } + PQclear(res); + PQfinish(conn); + + /* CopyDone cannot terminate an unfinished compressed frame. */ + conn = compression_connect(conninfo); + if (conn == NULL) + return 1; + if (start_copy(conn) || + send_compressed_copy_frame(conn, "unfinished\n", + strlen("unfinished\n"), false) || + expect_protocol_rejection(conn)) + goto done; + PQfinish(conn); + + /* A second frame cannot start before the COPY protocol boundary. */ + conn = compression_connect(conninfo); + if (conn == NULL) + return 1; + if (start_copy(conn) || + send_compressed_copy_frame(conn, "first\n", strlen("first\n"), true) || + send_compressed_copy_frame(conn, "second\n", strlen("second\n"), true) || + expect_protocol_rejection(conn)) + goto done; + + status = 0; + +done: + if (status != 0) + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return status; +} + +static int +test_compression_incompressible_copy(const char *conninfo) +{ + const int data_size = PQ_COMPRESSION_MAX_SEGMENT_SIZE + 1; + const size_t copy_size = (size_t) data_size + 27; + PGconn *conn = PQconnectdb(conninfo); + PGresult *res = NULL; + char *copy_data = NULL; + pg_prng_state state; + uint16 n16; + uint32 n32; + size_t offset; + int status = 1; + + if (PQstatus(conn) != CONNECTION_OK) + goto done; + copy_data = malloc(copy_size); + if (copy_data == NULL) + goto done; + + /* Binary COPY header, one bytea field, and the end-of-data marker. */ + memcpy(copy_data, "PGCOPY\n\377\r\n", 11); + memset(copy_data + 11, 0, 8); + n16 = pg_hton16(1); + memcpy(copy_data + 19, &n16, 2); + n32 = pg_hton32(data_size); + memcpy(copy_data + 21, &n32, 4); + pg_prng_seed(&state, 42); + for (offset = 0; offset < (size_t) data_size; offset += sizeof(uint64)) + { + uint64 value = pg_prng_uint64(&state); + size_t len = Min(sizeof(value), (size_t) data_size - offset); + + memcpy(copy_data + 25 + offset, &value, len); + } + n16 = pg_hton16(0xffff); + memcpy(copy_data + 25 + data_size, &n16, 2); + + res = PQexec(conn, "CREATE TEMP TABLE compression_copy_bytes (data bytea); " + "COPY compression_copy_bytes FROM STDIN (FORMAT binary)"); + if (PQresultStatus(res) != PGRES_COPY_IN) + goto done; + PQclear(res); + res = NULL; + if (PQputCopyData(conn, copy_data, (int) copy_size) != 1 || + PQputCopyEnd(conn, NULL) != 1 || + expect_copy_result(conn, PGRES_COMMAND_OK)) + goto done; + + res = PQexecParams(conn, "SELECT data FROM compression_copy_bytes", + 0, NULL, NULL, NULL, NULL, 1); + if (PQresultStatus(res) != PGRES_TUPLES_OK || + PQntuples(res) != 1 || PQnfields(res) != 1 || + PQgetisnull(res, 0, 0) || PQgetlength(res, 0, 0) != data_size || + memcmp(PQgetvalue(res, 0, 0), copy_data + 25, data_size) != 0) + goto done; + status = 0; + +done: + if (status != 0) + fprintf(stderr, "incompressible COPY round trip failed: %s", + PQerrorMessage(conn)); + PQclear(res); + free(copy_data); + PQfinish(conn); + return status; +} + +static int +test_compression_pipeline(const char *conninfo) +{ + const char *values[1]; + PGconn *conn = PQconnectdb(conninfo); + PGresult *res; + int status = 1; + + if (PQstatus(conn) != CONNECTION_OK || PQenterPipelineMode(conn) != 1) + goto done; + values[0] = "1"; + if (PQsendQueryParams(conn, + "SELECT $1::integer, repeat('pipeline result ', 1000)", + 1, NULL, values, NULL, NULL, 0) != 1) + goto done; + values[0] = "2"; + if (PQsendQueryParams(conn, + "SELECT $1::integer, repeat('pipeline result ', 1000)", + 1, NULL, values, NULL, NULL, 0) != 1 || + PQpipelineSync(conn) != 1) + goto done; + + res = PQgetResult(conn); + if (res == NULL || PQresultStatus(res) != PGRES_TUPLES_OK || + strcmp(PQgetvalue(res, 0, 0), "1") != 0) + goto done; + PQclear(res); + if (PQgetResult(conn) != NULL) + goto done; + res = PQgetResult(conn); + if (res == NULL || PQresultStatus(res) != PGRES_TUPLES_OK || + strcmp(PQgetvalue(res, 0, 0), "2") != 0) + goto done; + PQclear(res); + if (PQgetResult(conn) != NULL) + goto done; + res = PQgetResult(conn); + if (res == NULL || PQresultStatus(res) != PGRES_PIPELINE_SYNC) + goto done; + PQclear(res); + if (PQgetResult(conn) != NULL || PQexitPipelineMode(conn) != 1) + goto done; + + /* Creating the decompressor proves that a wrapper was received. */ + if (conn->compression_dctx == NULL) + goto done; + status = 0; + +done: + if (status != 0) + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return status; +} +#endif + +static int +test_compression_large_copy(const char *conninfo) +{ + const int data_size = 17 * 1024 * 1024; + PGconn *conn = PQconnectdb(conninfo); + PGresult *res; + char *data; + char small_data[1025]; + char trace_line[256]; + FILE *trace = NULL; + bool traced_compressed_data = false; + int status = 1; + + if (PQstatus(conn) != CONNECTION_OK) + goto fail; + res = PQexec(conn, "CREATE TEMP TABLE compression_copy (data text); " + "COPY compression_copy FROM STDIN"); + if (PQresultStatus(res) != PGRES_COPY_IN) + { + PQclear(res); + goto fail; + } + PQclear(res); + trace = tmpfile(); + if (trace == NULL) + goto fail; + PQtrace(conn, trace); + PQsetTraceFlags(conn, PQTRACE_SUPPRESS_TIMESTAMPS | PQTRACE_REGRESS_MODE); + + data = malloc(data_size + 1); + if (data == NULL) + goto fail; + memset(data, 'x', data_size); + data[data_size] = '\n'; + if (PQputCopyData(conn, data, data_size + 1) != 1 || + PQputCopyEnd(conn, NULL) != 1) + { + free(data); + goto fail; + } + free(data); + PQuntrace(conn); + rewind(trace); + while (fgets(trace_line, sizeof(trace_line), trace) != NULL) + { + if (strstr(trace_line, "CompressedData") != NULL) + traced_compressed_data = true; + if (strstr(trace_line, "mismatched message length") != NULL) + goto fail; + } + fclose(trace); + trace = NULL; + if (!traced_compressed_data) + goto fail; + while ((res = PQgetResult(conn)) != NULL) + { + if (PQresultStatus(res) != PGRES_COMMAND_OK) + { + PQclear(res); + goto fail; + } + PQclear(res); + } + + res = PQexec(conn, "SELECT count(*) = 1 AND " + "bool_and(data = repeat('x', 17 * 1024 * 1024)) " + "FROM compression_copy"); + if (PQresultStatus(res) != PGRES_TUPLES_OK || + strcmp(PQgetvalue(res, 0, 0), "t") != 0) + { + PQclear(res); + goto fail; + } + PQclear(res); + + /* Reuse the connection after large buffers were released at COPY end. */ + res = PQexec(conn, "COPY compression_copy FROM STDIN"); + if (PQresultStatus(res) != PGRES_COPY_IN) + { + PQclear(res); + goto fail; + } + PQclear(res); + memset(small_data, 'y', sizeof(small_data) - 1); + small_data[sizeof(small_data) - 1] = '\n'; + if (PQputCopyData(conn, small_data, sizeof(small_data)) != 1 || + PQputCopyData(conn, "short\n", 6) != 1 || + PQputCopyData(conn, small_data, sizeof(small_data)) != 1 || + PQputCopyEnd(conn, NULL) != 1) + goto fail; + while ((res = PQgetResult(conn)) != NULL) + { + if (PQresultStatus(res) != PGRES_COMMAND_OK) + { + PQclear(res); + goto fail; + } + PQclear(res); + } + res = PQexec(conn, "SELECT count(*) = 4 AND " + "count(*) FILTER (WHERE data = repeat('y', 1024)) = 2 AND " + "count(*) FILTER (WHERE data = 'short') = 1 " + "FROM compression_copy"); + if (PQresultStatus(res) == PGRES_TUPLES_OK && + strcmp(PQgetvalue(res, 0, 0), "t") == 0) + status = 0; + PQclear(res); + +fail: + if (trace != NULL) + { + PQuntrace(conn); + fclose(trace); + } + if (status != 0) + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return status; +} + +/* Exercise the message visibility and reset cases from the original series. */ +static int +test_compression_trace(const char *conninfo, const char *filename) +{ + static const struct + { + const char *sql; + ExecStatusType expected; + } commands[] = { + {"BEGIN", PGRES_COMMAND_OK}, + {"SELECT repeat('a', 80)", PGRES_TUPLES_OK}, + {"SET application_name = 'compressed trace'", PGRES_COMMAND_OK}, + {"DO $$ BEGIN RAISE NOTICE 'trace notice'; END $$", PGRES_COMMAND_OK}, + {"SELECT repeat('b', 80), 1 / i FROM generate_series(2, 0, -1) i", + PGRES_FATAL_ERROR}, + {"ROLLBACK", PGRES_COMMAND_OK}, + {"SELECT repeat('c', 80)", PGRES_TUPLES_OK} + }; + PGconn *conn = PQconnectdb(conninfo); + PGresult *res; + FILE *trace = NULL; + char *data; + int len; + int status = 1; + + if (PQstatus(conn) != CONNECTION_OK) + goto done; + trace = fopen(filename, "w"); + if (trace == NULL) + goto done; + PQtrace(conn, trace); + PQsetTraceFlags(conn, PQTRACE_SUPPRESS_TIMESTAMPS); + res = PQexec(conn, "SELECT repeat('trace payload ', 1000)"); + if (PQresultStatus(res) != PGRES_TUPLES_OK) + { + PQclear(res); + goto done; + } + PQclear(res); + PQsetTraceFlags(conn, PQTRACE_SUPPRESS_TIMESTAMPS | PQTRACE_REGRESS_MODE); + for (int i = 0; i < lengthof(commands); i++) + { + res = PQexec(conn, commands[i].sql); + if (PQresultStatus(res) != commands[i].expected) + { + PQclear(res); + goto done; + } + PQclear(res); +#ifdef USE_ZSTD + if (conn->compression_in_frame || conn->compression_frame_ended) + goto done; +#endif + } + res = PQexec(conn, + "COPY (SELECT repeat('d', 80) FROM generate_series(1, 3)) TO STDOUT"); + if (PQresultStatus(res) != PGRES_COPY_OUT) + { + PQclear(res); + goto done; + } + PQclear(res); + while ((len = PQgetCopyData(conn, &data, 0)) > 0) + PQfreemem(data); + if (len != -1) + goto done; + while ((res = PQgetResult(conn)) != NULL) + { + if (PQresultStatus(res) != PGRES_COMMAND_OK) + { + PQclear(res); + goto done; + } + PQclear(res); + } + status = 0; +done: + if (status != 0) + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQuntrace(conn); + if (trace) + fclose(trace); + PQfinish(conn); + return status; +} + int main(int argc, char *argv[]) { @@ -31,7 +708,29 @@ main(int argc, char *argv[]) print_ssl_library(); return 0; } + if (argc == 3 && !strcmp(argv[1], "--compression-reset")) + return test_compression_reset(argv[2]); + if (argc == 3 && !strcmp(argv[1], "--compression-large-copy")) + return test_compression_large_copy(argv[2]); + if (argc == 4 && !strcmp(argv[1], "--compression-trace")) + return test_compression_trace(argv[2], argv[3]); +#ifdef USE_ZSTD + if (argc == 3 && !strcmp(argv[1], "--compression-incompressible-copy")) + return test_compression_incompressible_copy(argv[2]); + if (argc == 3 && !strcmp(argv[1], "--compression-large-datarow")) + return test_compression_large_datarow(argv[2]); + if (argc == 3 && !strcmp(argv[1], "--compression-frame-boundaries")) + return test_compression_frame_boundaries(argv[2]); + if (argc == 3 && !strcmp(argv[1], "--compression-pipeline")) + return test_compression_pipeline(argv[2]); +#endif - printf("currently only --ssl is supported\n"); + printf("supported arguments are --ssl, --compression-reset CONNINFO, " + "--compression-large-copy CONNINFO, " + "--compression-incompressible-copy CONNINFO, " + "--compression-large-datarow CONNINFO, " + "--compression-frame-boundaries CONNINFO, " + "--compression-pipeline CONNINFO, and " + "--compression-trace CONNINFO FILE\n"); return 1; } diff --git a/src/interfaces/libpq/test/meson.build b/src/interfaces/libpq/test/meson.build index e203486615c..2325eb6dcd7 100644 --- a/src/interfaces/libpq/test/meson.build +++ b/src/interfaces/libpq/test/meson.build @@ -25,6 +25,11 @@ libpq_testclient_sources = files( 'libpq_testclient.c', ) +libpq_testclient_deps = [frontend_no_fe_utils_code, libpq] +if zstd.found() + libpq_testclient_deps += zstd +endif + if host_system == 'windows' libpq_testclient_sources += rc_bin_gen.process(win32ver_rc, extra_args: [ '--NAME', 'libpq_testclient', @@ -33,7 +38,8 @@ endif libpq_test_deps += executable('libpq_testclient', libpq_testclient_sources, - dependencies: [frontend_no_fe_utils_code, libpq], + include_directories: [libpq_inc], + dependencies: libpq_testclient_deps, kwargs: default_bin_args + { 'install': false, } diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list index c42dd48bb23..f350e549631 100644 --- a/src/tools/pgindent/typedefs.list +++ b/src/tools/pgindent/typedefs.list @@ -2428,6 +2428,7 @@ ProjectSetState ProjectionInfo ProjectionPath PromptInterruptContext +ProtocolCompression ProtocolVersion PrsStorage PruneFreezeParams @@ -3540,6 +3541,7 @@ ZSTD_CCtx ZSTD_CStream ZSTD_DCtx ZSTD_DStream +ZSTD_EndDirective ZSTD_cParameter ZSTD_inBuffer ZSTD_outBuffer -- That's all, folks. May the source be with you.