From 5ae4d9bd03f52efce46cd07ead1b35939939ad13 Mon Sep 17 00:00:00 2001 From: Andrey Borodin Date: Sun, 30 Aug 2026 14:30:50 +0500 Subject: [PATCH v1] Add protocol compression Add optional Zstandard compression for DataRow and CopyData protocol messages. Negotiate it with the _pq_.compression startup option, a server GUC, and the libpq compression connection option. Compression starts after the first ReadyForQuery. Carry eligible messages in bounded segments of persistent streams with a 64 KiB history. Reset frames at ReadyForQuery and COPY boundaries, enforce wrapper, output, frame, and inner-message boundaries, and leave queries, bind parameters, and control messages uncompressed. Compress server results and COPY output, plus frontend CopyData during COPY FROM. Add tests for negotiation, query and COPY equivalence, pg_restore, incremental backup, large segmented messages, errors, reconnection, pipeline mode, tracing, and frame boundaries. --- doc/src/sgml/config.sgml | 18 + doc/src/sgml/libpq.sgml | 44 ++ doc/src/sgml/protocol.sgml | 71 +- meson.build | 1 + src/backend/commands/copyfromparse.c | 22 +- src/backend/libpq/pqcomm.c | 672 +++++++++++++++++- src/backend/replication/walsender.c | 26 +- src/backend/tcop/backend_startup.c | 21 +- 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 | 1 + src/include/libpq/libpq.h | 13 + src/include/libpq/protocol.h | 5 + src/interfaces/libpq/Makefile | 1 + src/interfaces/libpq/fe-connect.c | 34 + src/interfaces/libpq/fe-exec.c | 23 + src/interfaces/libpq/fe-protocol3.c | 539 ++++++++++++++ src/interfaces/libpq/fe-trace.c | 5 + src/interfaces/libpq/libpq-int.h | 21 + src/interfaces/libpq/meson.build | 1 + src/interfaces/libpq/t/007_compression.pl | 182 +++++ src/interfaces/libpq/test/libpq_testclient.c | 502 ++++++++++++- src/interfaces/libpq/test/meson.build | 8 +- 24 files changed, 2227 insertions(+), 18 deletions(-) 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 0165eb9ec02..5a9b14c9fdb 100644 --- a/doc/src/sgml/config.sgml +++ b/doc/src/sgml/config.sgml @@ -963,6 +963,24 @@ 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. + + + diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml index 123e7f03902..4e8b21c6b95 100644 --- a/doc/src/sgml/libpq.sgml +++ b/doc/src/sgml/libpq.sgml @@ -2251,6 +2251,40 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname + + compression + + + Enables compression of data-bearing messages sent by the server. + 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 @@ -9330,6 +9364,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 49f81676712..02c574225a7 100644 --- a/doc/src/sgml/protocol.sgml +++ b/doc/src/sgml/protocol.sgml @@ -346,9 +346,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 + . @@ -4656,6 +4661,66 @@ 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 + is limited to 17 MiB and may produce at most 16 MiB of uncompressed data. + 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 f4cde249242..f1369a8e879 100644 --- a/meson.build +++ b/meson.build @@ -3619,6 +3619,7 @@ frontend_no_fe_utils_code = declare_dependency( # Dependencies both for static and shared libpq libpq_deps += [ thread_dep, + zstd, gssapi, ldap_r, diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c index 37750cca13a..e1026530798 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/pqcomm.c b/src/backend/libpq/pqcomm.c index 3704d121003..d95cd0356b6 100644 --- a/src/backend/libpq/pqcomm.c +++ b/src/backend/libpq/pqcomm.c @@ -69,7 +69,12 @@ #include #endif +#ifdef USE_ZSTD +#include +#endif + #include "common/ip.h" +#include "common/int.h" #include "libpq/libpq.h" #include "miscadmin.h" #include "port/pg_bswap.h" @@ -132,6 +137,25 @@ static int PqRecvLength; /* End of data available in PqRecvBuffer */ static bool PqCommBusy; /* busy sending data to the client */ static bool PqCommReadingMsg; /* in the middle of reading a message */ +#ifdef USE_ZSTD +/* Experimental server-to-client protocol compression. */ +static bool PqCompressionStarted; +static bool PqCompressionNegotiated; +static bool PqCompressionActive; +static bool PqCompressionFrameStarted; +static size_t PqCompressionSmallBytes; +static StringInfoData PqCompressionInput; +static StringInfoData PqCompressionOutput; +static StringInfoData PqDecompressionBuffer; +static bool PqDecompressionBufferInitialized; +static bool PqDecompressionFrameStarted; +static bool PqDecompressionFrameEnded; +static ZSTD_CCtx *PqCompressionContext; +static ZSTD_DCtx *PqDecompressionContext; +#endif + +int protocol_compression = PROTOCOL_COMPRESSION_OFF; + /* Internal functions */ static void socket_comm_reset(void); @@ -142,6 +166,14 @@ static int socket_flush_if_writable(void); static bool socket_is_send_pending(void); static int socket_putmessage(char msgtype, const char *s, size_t len); static void socket_putmessage_noblock(char msgtype, const char *s, size_t len); +#ifdef USE_ZSTD +static int socket_compression_flush(ZSTD_EndDirective directive); +static int socket_compression_init(void); +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 socket_decompression_reset(void); +#endif static inline int internal_putbytes(const void *b, size_t len); static inline int internal_flush(void); static pg_noinline int internal_flush_buffer(const char *buf, size_t *start, @@ -1270,6 +1302,200 @@ pq_getmessage(StringInfo s, int maxlen) return 0; } +#ifdef USE_ZSTD +/* Check a physical frontend message against the compression frame state. */ +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"))); + PqDecompressionFrameStarted = false; + PqDecompressionFrameEnded = false; + } +} + +/* Read one client-to-server CompressedData wrapper. */ +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(&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, 16); + if (ZSTD_isError(rc)) + ereport(FATAL, + (errcode(ERRCODE_INTERNAL_ERROR), + errmsg("could not configure Zstandard decompression context: %s", + ZSTD_getErrorName(rc)))); + } + + resetStringInfo(&PqDecompressionBuffer); + for (;;) + { + StringInfoData compressed; + ZSTD_inBuffer input; + size_t segment_size = 0; + + initStringInfo(&compressed); + if (pq_getmessage(&compressed, PQ_COMPRESSION_MAX_WRAPPER_SIZE + 4)) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("invalid compressed protocol message"))); + if (compressed.len == 0) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("compressed protocol message is empty"))); + + input.src = compressed.data; + input.size = compressed.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; + } + pfree(compressed.data); + 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(); +} +#endif + static inline int internal_putbytes(const void *b, size_t len) @@ -1330,6 +1556,13 @@ socket_flush(void) if (PqCommBusy) return 0; PqCommBusy = true; +#ifdef USE_ZSTD + if (PqCompressionStarted && socket_compression_flush(ZSTD_e_flush)) + { + PqCommBusy = false; + return EOF; + } +#endif socket_set_nonblocking(false); res = internal_flush(); PqCommBusy = false; @@ -1435,7 +1668,11 @@ socket_flush_if_writable(void) int res; /* Quick exit if nothing to do */ - if (PqSendPointer == PqSendStart) + if (PqSendPointer == PqSendStart +#ifdef USE_ZSTD + && (!PqCompressionStarted || PqCompressionInput.len == 0) +#endif + ) return 0; /* No-op if reentrant call */ @@ -1446,6 +1683,13 @@ socket_flush_if_writable(void) socket_set_nonblocking(true); PqCommBusy = true; +#ifdef USE_ZSTD + if (PqCompressionStarted && socket_compression_flush(ZSTD_e_flush)) + { + PqCommBusy = false; + return EOF; + } +#endif res = internal_flush(); PqCommBusy = false; return res; @@ -1458,6 +1702,10 @@ socket_flush_if_writable(void) static bool socket_is_send_pending(void) { +#ifdef USE_ZSTD + if (PqCompressionStarted && PqCompressionInput.len > 0) + return true; +#endif return (PqSendStart < PqSendPointer); } @@ -1466,6 +1714,312 @@ socket_is_send_pending(void) * -------------------------------- */ +#ifdef USE_ZSTD +void +pq_enable_protocol_compression(void) +{ + if (PqCompressionNegotiated) + ereport(FATAL, + (errcode(ERRCODE_PROTOCOL_VIOLATION), + errmsg("protocol compression option specified more than once"))); + PqCompressionNegotiated = true; +} + +static int +socket_compression_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, 16); + 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; +} + +/* Initialize a compression buffer in TopMemoryContext, without throwing. */ +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; +} + +/* + * Make room for size more bytes in a compression buffer, without throwing. + * + * The compression code runs with PqCommBusy set, so an error raised here would + * be reported into a message stream that cannot carry it, and would leave + * PqCommBusy set for the rest of the session. Report a failure to the caller + * instead, the way the rest of the send path does. + */ +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; +} + +/* + * Compression buffers up to this size are kept between messages; a larger one + * was grown by an occasional large message and is given back at the next + * protocol boundary. The steady state needs the flush output buffer to hold + * ZSTD_CStreamOutSize() bytes, so keep well above that. + */ +#define PQ_COMPRESSION_BUFFER_KEEP_SIZE (1024 * 1024) + +/* + * Give back a compression buffer that one large message has grown. + * + * Called where the buffer holds nothing live, so that a single large result + * does not keep tens of megabytes for the life of the backend. + */ +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'; +} + +/* Start the frontend decompressor from a clean protocol boundary. */ +static void +socket_decompression_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(&PqDecompressionBuffer); + PqDecompressionFrameStarted = false; + PqDecompressionFrameEnded = false; +} + +/* Upper bound on the wire size of one flushed segment of input_len bytes. */ +static inline Size +compression_segment_bound(Size input_len) +{ + return ZSTD_compressBound(input_len) + ZSTD_CStreamOutSize() + 64 + 5; +} + +/* Emit one completely flushed segment of the persistent stream. */ +static int +socket_compression_flush(ZSTD_EndDirective directive) +{ + size_t output_size; + ZSTD_inBuffer input; + ZSTD_outBuffer out; + size_t result; + uint32 n32; + + /* Do not create an empty frame for a result that was kept uncompressed. */ + if (!PqCompressionFrameStarted) + return 0; + + output_size = ZSTD_compressBound(PqCompressionInput.len) + + ZSTD_CStreamOutSize() + 64; + 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) + { + char msgtype = PqMsg_CompressedData; + + if (internal_putbytes(&msgtype, 1)) + goto fail; + n32 = pg_hton32((uint32) (out.pos + 4)); + if (internal_putbytes(&n32, 4)) + goto fail; + if (internal_putbytes(PqCompressionOutput.data, out.pos)) + goto fail; + } + + 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; +} + +/* Append logical protocol bytes, splitting them at wrapper boundaries. */ +static int +socket_compression_append(const void *data, Size len) +{ + const char *ptr = data; + + while (len > 0) + { + Size available = PQ_COMPRESSION_MAX_SEGMENT_SIZE - + PqCompressionInput.len; + Size part = Min(len, available); + + if (part == 0) + { + if (socket_compression_flush(ZSTD_e_flush)) + 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_MAX_SEGMENT_SIZE && + socket_compression_flush(ZSTD_e_flush)) + return EOF; + } + return 0; +} +#endif + /* -------------------------------- * socket_putmessage - send a normal message (suppressed in COPY OUT mode) @@ -1495,6 +2049,48 @@ socket_putmessage(char msgtype, const char *s, size_t len) if (PqCommBusy) return 0; PqCommBusy = true; + +#ifdef USE_ZSTD + if (PqCompressionStarted && + (msgtype == PqMsg_DataRow || msgtype == PqMsg_CopyData)) + { + size_t message_size = len + 5; + + if (!PqCompressionActive && len + 4 < 60) + { + PqCompressionSmallBytes += message_size; + if (PqCompressionSmallBytes >= 1024) + PqCompressionActive = true; + } + else + { + if (PqCompressionContext == NULL) + { + if (socket_compression_init()) + goto fail; + } + PqCompressionActive = true; + PqCompressionFrameStarted = true; + n32 = pg_hton32((uint32) (len + 4)); + if (socket_compression_append(&msgtype, 1) || + socket_compression_append(&n32, 4) || + socket_compression_append(s, len)) + goto fail; + if (PqCompressionInput.len >= PQ_SEND_BUFFER_SIZE && + socket_compression_flush(ZSTD_e_flush)) + goto fail; + PqCommBusy = false; + return 0; + } + } + + if (PqCompressionStarted && + msgtype != PqMsg_DataRow && msgtype != PqMsg_CopyData && + socket_compression_flush(msgtype == PqMsg_ReadyForQuery ? + ZSTD_e_end : ZSTD_e_flush)) + goto fail; +#endif + if (internal_putbytes(&msgtype, 1)) goto fail; @@ -1505,6 +2101,22 @@ socket_putmessage(char msgtype, const char *s, size_t len) if (internal_putbytes(s, len)) goto fail; PqCommBusy = false; +#ifdef USE_ZSTD + if (PqCompressionNegotiated && msgtype == PqMsg_ReadyForQuery) + { + PqCompressionStarted = true; + PqCompressionActive = false; + PqCompressionFrameStarted = false; + PqCompressionSmallBytes = 0; + socket_decompression_reset(); + + /* The frame ended above, so none of the buffers holds anything now. */ + Assert(PqCompressionInput.len == 0); + compression_buffer_release(&PqCompressionInput); + compression_buffer_release(&PqCompressionOutput); + compression_buffer_release(&PqDecompressionBuffer); + } +#endif return 0; fail: @@ -1522,21 +2134,69 @@ static void socket_putmessage_noblock(char msgtype, const char *s, size_t len) { int res PG_USED_FOR_ASSERTS_ONLY; - int required; + Size required; /* * Ensure we have enough space in the output buffer for the message header * as well as the message itself. */ - required = PqSendPointer + 1 + 4 + len; + if (pg_add_size_overflow(PqSendPointer, 5, &required) || + pg_add_size_overflow(required, len, &required)) + ereport(ERROR, + (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), + errmsg("message is too long"))); +#ifdef USE_ZSTD + if (PqCompressionStarted) + { + Size pending = PqCompressionInput.len; + Size message_size; + Size compression_space = 0; + + if (pg_add_size_overflow(len, 5, &message_size)) + ereport(ERROR, + (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), + errmsg("message is too long"))); + + if (msgtype == PqMsg_DataRow || msgtype == PqMsg_CopyData) + { + Size remaining = message_size; + + while (remaining > PQ_COMPRESSION_MAX_SEGMENT_SIZE - pending) + { + Size part = PQ_COMPRESSION_MAX_SEGMENT_SIZE - pending; + + remaining -= part; + if (pg_add_size_overflow(compression_space, + compression_segment_bound(PQ_COMPRESSION_MAX_SEGMENT_SIZE), + &compression_space)) + ereport(ERROR, + (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), + errmsg("message is too long"))); + pending = 0; + } + pending += remaining; + } + if ((pending > 0 && + pg_add_size_overflow(compression_space, + compression_segment_bound(pending), + &compression_space)) || + pg_add_size_overflow(required, compression_space, &required)) + ereport(ERROR, + (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), + errmsg("message is too long"))); + } +#endif + if (!AllocSizeIsValid(required)) + ereport(ERROR, + (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), + errmsg("message is too long"))); if (required > PqSendBufferSize) { PqSendBuffer = repalloc(PqSendBuffer, required); - PqSendBufferSize = required; + PqSendBufferSize = (int) required; } res = socket_putmessage(msgtype, s, len); - Assert(res == 0); /* should not fail when the message fits in - * buffer */ + Assert(res == 0 || ClientConnectionLost); } /* -------------------------------- diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c index c65dd324325..5c6fb3510e8 100644 --- a/src/backend/replication/walsender.c +++ b/src/backend/replication/walsender.c @@ -786,9 +786,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) @@ -796,6 +800,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: @@ -817,7 +839,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..30e384cbc12 100644 --- a/src/backend/tcop/backend_startup.c +++ b/src/backend/tcop/backend_startup.c @@ -808,12 +808,29 @@ retry: valptr), errhint("Valid values are: \"false\", 0, \"true\", 1, \"database\"."))); } + else if (strcmp(nameptr, "_pq_.compression") == 0) + { +#ifdef USE_ZSTD + if (strcmp(valptr, "zstd") != 0) + ereport(FATAL, + (errcode(ERRCODE_INVALID_PARAMETER_VALUE), + errmsg("invalid value for protocol option \"%s\": \"%s\"", + nameptr, valptr))); + if (protocol_compression == PROTOCOL_COMPRESSION_ZSTD) + pq_enable_protocol_compression(); + else + unrecognized_protocol_options = + lappend(unrecognized_protocol_options, pstrdup(nameptr)); +#else + unrecognized_protocol_options = + lappend(unrecognized_protocol_options, pstrdup(nameptr)); +#endif + } 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 - * defined. + * protocol-level option. */ unrecognized_protocol_options = lappend(unrecognized_protocol_options, pstrdup(nameptr)); 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 3c5e16ad1e7..68f6261fb5f 100644 --- a/src/backend/utils/misc/guc_parameters.dat +++ b/src/backend/utils/misc/guc_parameters.dat @@ -2429,6 +2429,13 @@ check_hook => 'check_primary_slot_name', }, +{ name => 'protocol_compression', type => 'enum', context => 'PGC_SIGHUP', group => 'CONN_AUTH_SETTINGS', + short_desc => 'Allows clients to negotiate protocol compression.', + 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 c6d9b2a6f89..ab37feacc08 100644 --- a/src/backend/utils/misc/guc_tables.c +++ b/src/backend/utils/misc/guc_tables.c @@ -496,6 +496,14 @@ static const struct config_enum_entry wal_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 file_copy_method_options[] = { {"copy", FILE_COPY_METHOD_COPY, false}, #if defined(HAVE_COPYFILE) && defined(COPYFILE_CLONE_FORCE) || defined(HAVE_COPY_FILE_RANGE) diff --git a/src/backend/utils/misc/postgresql.conf.sample b/src/backend/utils/misc/postgresql.conf.sample index e759f06b50f..c83a76bc422 100644 --- a/src/backend/utils/misc/postgresql.conf.sample +++ b/src/backend/utils/misc/postgresql.conf.sample @@ -76,6 +76,7 @@ # (change requires restart) #bonjour_name = '' # defaults to the computer name # (change requires restart) +#protocol_compression = off # off, zstd # - TCP settings - # see "man tcp" for details diff --git a/src/include/libpq/libpq.h b/src/include/libpq/libpq.h index d15073a0a93..f5cf0f12b85 100644 --- a/src/include/libpq/libpq.h +++ b/src/include/libpq/libpq.h @@ -33,6 +33,12 @@ typedef struct WaitEventSet WaitEventSet; #define PQ_SMALL_MESSAGE_LIMIT 10000 #define PQ_LARGE_MESSAGE_LIMIT (MaxAllocSize - 1) +typedef enum ProtocolCompressionMethod +{ + PROTOCOL_COMPRESSION_OFF, + PROTOCOL_COMPRESSION_ZSTD +} ProtocolCompressionMethod; + typedef struct { void (*comm_reset) (void); @@ -99,6 +105,13 @@ extern ssize_t secure_write(Port *port, const void *ptr, size_t len); extern ssize_t secure_raw_read(Port *port, void *ptr, size_t len); extern ssize_t secure_raw_write(Port *port, const void *ptr, size_t len); +#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 +extern PGDLLIMPORT int protocol_compression; + /* * declarations for variables defined in be-secure.c */ diff --git a/src/include/libpq/protocol.h b/src/include/libpq/protocol.h index eae8f0e7238..7f4f4e059e2 100644 --- a/src/include/libpq/protocol.h +++ b/src/include/libpq/protocol.h @@ -57,6 +57,11 @@ #define PqMsg_PortalSuspended 's' #define PqMsg_ParameterDescription 't' #define PqMsg_NegotiateProtocolVersion 'v' +#define PqMsg_CompressedData 'z' + +/* Resource limits for one CompressedData payload and its decoded segment. */ +#define PQ_COMPRESSION_MAX_SEGMENT_SIZE (16 * 1024 * 1024) +#define PQ_COMPRESSION_MAX_WRAPPER_SIZE (17 * 1024 * 1024) /* These are the codes sent by both the frontend and backend. */ diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 0963995eed4..414db507c40 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -86,6 +86,7 @@ endif # instead link with -lpgcommon_shlib and -lpgport_shlib, to get object files # that are built correctly for use in a shlib. SHLIB_LINK_INTERNAL = -lpgcommon_shlib -lpgport_shlib +SHLIB_LINK += $(filter -lzstd, $(LIBS)) 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) else diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c index ee398f13998..b09af588953 100644 --- a/src/interfaces/libpq/fe-connect.c +++ b/src/interfaces/libpq/fe-connect.c @@ -347,6 +347,10 @@ static const internalPQconninfoOption PQconninfoOptions[] = { "Max-Protocol-Version", "", 6, /* sizeof("latest") = 6 */ offsetof(struct pg_conn, max_protocol_version)}, + {"compression", "PGCOMPRESSION", "off", NULL, + "Protocol-Compression", "", 7, /* sizeof("prefer") == 7 */ + offsetof(struct pg_conn, compression)}, + {"ssl_min_protocol_version", "PGSSLMINPROTOCOLVERSION", "TLSv1.2", NULL, "SSL-Minimum-Protocol-Version", "", 8, /* sizeof("TLSv1.x") == 8 */ offsetof(struct pg_conn, ssl_min_protocol_version)}, @@ -2166,6 +2170,27 @@ pqConnectOptions2(PGconn *conn) return false; } + if (conn->compression && strcmp(conn->compression, "off") != 0) + { + if (strcmp(conn->compression, "zstd") != 0 && + strcmp(conn->compression, "prefer") != 0) + { + conn->status = CONNECTION_BAD; + libpq_append_conn_error(conn, "invalid %s value: \"%s\"", + "compression", conn->compression); + return false; + } +#ifndef USE_ZSTD + if (strcmp(conn->compression, "zstd") == 0) + { + conn->status = CONNECTION_BAD; + libpq_append_conn_error(conn, + "compression method \"zstd\" is not supported by this build"); + return false; + } +#endif + } + /* * Resolve special "auto" client_encoding from the locale */ @@ -5152,6 +5177,7 @@ freePGconn(PGconn *conn) free(conn->gssdelegation); free(conn->min_protocol_version); free(conn->max_protocol_version); + free(conn->compression); free(conn->ssl_min_protocol_version); free(conn->ssl_max_protocol_version); free(conn->target_session_attrs); @@ -5174,6 +5200,9 @@ freePGconn(PGconn *conn) release_conn_addrinfo(conn); free(conn->scram_client_key_binary); free(conn->scram_server_key_binary); +#ifdef USE_ZSTD + pqCompressionReset(conn); +#endif /* if this is a cancel connection, be_cancel_key may still be allocated */ free(conn->be_cancel_key); free(conn->inBuffer); @@ -5347,6 +5376,11 @@ pqClosePGconn(PGconn *conn) /* Reset all state obtained from server, too */ pqDropServerData(conn); +#ifdef USE_ZSTD + pqCompressionReset(conn); +#else + conn->compression_rejected = false; +#endif } /* 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-protocol3.c b/src/interfaces/libpq/fe-protocol3.c index 79099301abd..f1f9103fe33 100644 --- a/src/interfaces/libpq/fe-protocol3.c +++ b/src/interfaces/libpq/fe-protocol3.c @@ -18,6 +18,10 @@ #include #include +#ifdef USE_ZSTD +#include +#endif + #ifdef WIN32 #include "win32.h" #else @@ -37,6 +41,7 @@ */ #define VALID_LONG_MESSAGE_TYPE(id) \ ((id) == PqMsg_CopyData || \ + (id) == PqMsg_CompressedData || \ (id) == PqMsg_DataRow || \ (id) == PqMsg_ErrorResponse || \ (id) == PqMsg_FunctionCallResponse || \ @@ -61,6 +66,409 @@ static void reportErrorPosition(PQExpBuffer msg, const char *query, static size_t build_startup_packet(const PGconn *conn, char *packet, const PQEnvironmentOption *options); +#ifdef USE_ZSTD + +/* + * Compression buffers up to this size are kept between messages; a larger one + * was grown by an occasional large message and is given back at the next + * protocol boundary. + */ +#define PQ_COMPRESSION_BUFFER_KEEP_SIZE (1024 * 1024) + +static int +pqCompressionInitBuffers(PGconn *conn) +{ + if (conn->compression_buffers_initialized) + return 0; + + initPQExpBuffer(&conn->compression_buffer); + initPQExpBuffer(&conn->compression_output_buffer); + if (PQExpBufferBroken(&conn->compression_buffer) || + PQExpBufferBroken(&conn->compression_output_buffer)) + { + termPQExpBuffer(&conn->compression_buffer); + termPQExpBuffer(&conn->compression_output_buffer); + libpq_append_conn_error(conn, "out of memory"); + return 1; + } + + conn->compression_buffers_initialized = true; + return 0; +} + +/* + * Give back the compression buffers that one large message has grown. + * + * Called where they hold nothing live, so that a single large result does not + * keep tens of megabytes for the life of the connection. + */ +static void +pqCompressionReleaseBuffers(PGconn *conn) +{ + char *data; + + if (!conn->compression_buffers_initialized) + return; + if (conn->compression_buffer.maxlen > PQ_COMPRESSION_BUFFER_KEEP_SIZE) + { + data = realloc(conn->compression_buffer.data, INITIAL_EXPBUFFER_SIZE); + if (data != NULL) + { + conn->compression_buffer.data = data; + conn->compression_buffer.maxlen = INITIAL_EXPBUFFER_SIZE; + resetPQExpBuffer(&conn->compression_buffer); + } + } + if (conn->compression_output_buffer.maxlen > PQ_COMPRESSION_BUFFER_KEEP_SIZE) + { + data = realloc(conn->compression_output_buffer.data, + INITIAL_EXPBUFFER_SIZE); + if (data != NULL) + { + conn->compression_output_buffer.data = data; + conn->compression_output_buffer.maxlen = INITIAL_EXPBUFFER_SIZE; + resetPQExpBuffer(&conn->compression_output_buffer); + } + } +} + +void +pqCompressionReset(PGconn *conn) +{ + 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; + } + if (conn->compression_buffers_initialized) + { + termPQExpBuffer(&conn->compression_buffer); + termPQExpBuffer(&conn->compression_output_buffer); + conn->compression_buffers_initialized = false; + } + conn->compression_in_frame = false; + conn->compression_frame_ended = false; + conn->compression_ready = false; + conn->compression_copy_started = false; + conn->compression_rejected = false; +} + +/* Replace one CompressedData wrapper with the ordinary messages it contains. */ +static int +pqDecompressData(PGconn *conn, int msgLength) +{ + int outer_size = msgLength + 5; + int tail_size; + int new_end; + int position; + size_t segment_size = 0; + ZSTD_inBuffer input; + ZSTD_DCtx *dctx; + size_t result = 1; + + if (msgLength <= 0 || msgLength > PQ_COMPRESSION_MAX_WRAPPER_SIZE) + { + libpq_append_conn_error(conn, "invalid compressed protocol message size"); + return 1; + } + + if (conn->compression_dctx == NULL) + { + size_t rc; + + dctx = ZSTD_createDCtx(); + if (dctx == NULL) + { + libpq_append_conn_error(conn, "out of memory"); + return 1; + } + if (pqCompressionInitBuffers(conn)) + { + ZSTD_freeDCtx(dctx); + return 1; + } + rc = ZSTD_DCtx_setParameter(dctx, ZSTD_d_windowLogMax, 16); + if (ZSTD_isError(rc)) + { + ZSTD_freeDCtx(dctx); + libpq_append_conn_error(conn, + "could not configure Zstandard decompressor: %s", + ZSTD_getErrorName(rc)); + return 1; + } + conn->compression_dctx = dctx; + } + dctx = (ZSTD_DCtx *) conn->compression_dctx; + if (conn->compression_frame_ended) + { + libpq_append_conn_error(conn, + "received data after the compressed protocol stream ended"); + return 1; + } + + input.src = conn->inBuffer + conn->inCursor; + input.size = msgLength; + 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 (!enlargePQExpBuffer(&conn->compression_buffer, output_size)) + { + libpq_append_conn_error(conn, "out of memory"); + return 1; + } + out.dst = conn->compression_buffer.data + conn->compression_buffer.len; + out.size = output_size; + out.pos = 0; + + result = ZSTD_decompressStream(dctx, &out, &input); + if (ZSTD_isError(result)) + { + libpq_append_conn_error(conn, "Zstandard decompression failed: %s", + ZSTD_getErrorName(result)); + return 1; + } + if (input.pos == old_input_pos && out.pos == 0) + { + if (input.pos == input.size) + break; + libpq_append_conn_error(conn, + "compressed protocol message has invalid size"); + return 1; + } + conn->compression_buffer.len += out.pos; + conn->compression_buffer.data[conn->compression_buffer.len] = '\0'; + segment_size += out.pos; + if (segment_size > PQ_COMPRESSION_MAX_SEGMENT_SIZE) + { + libpq_append_conn_error(conn, + "compressed protocol message is too large"); + return 1; + } + if (result == 0) + { + if (input.pos != input.size) + { + libpq_append_conn_error(conn, + "compressed protocol message contains multiple frames"); + return 1; + } + break; + } + if (input.pos == input.size && out.pos < out.size) + break; + } + conn->compression_in_frame = (result != 0); + conn->compression_frame_ended = (result == 0); + + /* Validate the complete prefix of the logical protocol stream. */ + position = 0; + while (conn->compression_buffer.len - position >= 5) + { + uint32 message_length; + + if (conn->compression_buffer.data[position] != PqMsg_DataRow && + conn->compression_buffer.data[position] != PqMsg_CopyData) + goto invalid_contents; + memcpy(&message_length, + conn->compression_buffer.data + position + 1, 4); + message_length = pg_ntoh32(message_length); + if (message_length < 4 || message_length > INT_MAX) + goto invalid_contents; + if (message_length > conn->compression_buffer.len - position - 1) + break; + position += message_length + 1; + } + if (result == 0 && position != conn->compression_buffer.len) + goto invalid_contents; + + tail_size = conn->inEnd - conn->inStart - outer_size; + + /* Keep an incomplete inner message until another wrapper arrives. */ + if (position != conn->compression_buffer.len) + position = 0; + new_end = conn->inEnd - outer_size + position; + if (pqCheckInBufferSpace(new_end, conn)) + return 1; + new_end = conn->inEnd - outer_size + position; + memmove(conn->inBuffer + conn->inStart + position, + conn->inBuffer + conn->inStart + outer_size, tail_size); + if (position > 0) + { + memcpy(conn->inBuffer + conn->inStart, conn->compression_buffer.data, + position); + resetPQExpBuffer(&conn->compression_buffer); + } + conn->inEnd = new_end; + conn->inCursor = conn->inStart; + return 0; + +invalid_contents: + libpq_append_conn_error(conn, + "compressed protocol message contains invalid messages"); + return 1; +} + +/* Write one frontend message as one or more compressed stream segments. */ +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, 16); + 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) + + ZSTD_CStreamOutSize() + 64; + 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; + } + 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; + } 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)); +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 + /* * parseInput: if appropriate, parse input data from backend @@ -104,6 +512,24 @@ pqParseInput3(PGconn *conn) handleSyncLoss(conn, id, msgLength); return; } +#ifdef USE_ZSTD + if (id == PqMsg_CompressedData && + msgLength - 4 > PQ_COMPRESSION_MAX_WRAPPER_SIZE) + { + libpq_append_conn_error(conn, + "compressed protocol message is too large"); + handleFatalError(conn); + return; + } + if (conn->compression_buffers_initialized && + conn->compression_buffer.len > 0 && id != PqMsg_CompressedData) + { + libpq_append_conn_error(conn, + "compressed protocol message was interrupted"); + handleFatalError(conn); + return; + } +#endif /* * Can't process if message body isn't all here yet. @@ -134,6 +560,43 @@ pqParseInput3(PGconn *conn) return; } +#ifdef USE_ZSTD + if (id == PqMsg_CompressedData) + { + if (!conn->compression_ready) + { + libpq_append_conn_error(conn, + "received compressed data without negotiated compression"); + handleFatalError(conn); + return; + } + if (pqDecompressData(conn, msgLength)) + { + handleFatalError(conn); + return; + } + continue; + } + if (id == PqMsg_ReadyForQuery && conn->compression_in_frame) + { + libpq_append_conn_error(conn, + "compressed protocol stream was not terminated before ReadyForQuery"); + handleFatalError(conn); + return; + } + if (id == PqMsg_ReadyForQuery) + { + 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; + pqCompressionReleaseBuffers(conn); + } +#endif + /* * NOTIFY and NOTICE messages can happen in any state; always process * them right away. @@ -228,6 +691,12 @@ pqParseInput3(PGconn *conn) case PqMsg_ReadyForQuery: if (getReadyForQuery(conn)) return; +#ifdef USE_ZSTD + if (conn->compression && !conn->compression_rejected && + (strcmp(conn->compression, "zstd") == 0 || + strcmp(conn->compression, "prefer") == 0)) + conn->compression_ready = true; +#endif if (conn->pipelineStatus != PQ_PIPELINE_OFF) { conn->result = PQmakeEmptyPGresult(conn, @@ -1447,6 +1916,7 @@ pqGetNegotiateProtocolVersion3(PGconn *conn) int num; bool found_test_protocol_negotiation; bool expect_test_protocol_negotiation; + bool requested_compression; /* * During 19beta only, if protocol grease is in use, assume that it's the @@ -1527,6 +1997,13 @@ pqGetNegotiateProtocolVersion3(PGconn *conn) */ found_test_protocol_negotiation = false; expect_test_protocol_negotiation = (conn->max_pversion == PG_PROTOCOL_GREASE); +#ifdef USE_ZSTD + requested_compression = conn->compression && + (strcmp(conn->compression, "zstd") == 0 || + strcmp(conn->compression, "prefer") == 0); +#else + requested_compression = false; +#endif for (int i = 0; i < num; i++) { @@ -1546,6 +2023,19 @@ pqGetNegotiateProtocolVersion3(PGconn *conn) { found_test_protocol_negotiation = true; } + else if (requested_compression && + strcmp(conn->workBuffer.data, "_pq_.compression") == 0) + { + if (strcmp(conn->compression, "prefer") == 0) + conn->compression_rejected = true; + else + { + libpq_append_conn_error(conn, + "server does not support protocol compression method \"zstd\""); + need_grease_info = false; + goto failure; + } + } else { libpq_append_conn_error(conn, "received invalid protocol negotiation message: server reported an unsupported parameter that was not requested (\"%s\")", @@ -1865,6 +2355,24 @@ getCopyDataMessage(PGconn *conn) handleSyncLoss(conn, id, msgLength); return -2; } +#ifdef USE_ZSTD + if (id == PqMsg_CompressedData && + msgLength - 4 > PQ_COMPRESSION_MAX_WRAPPER_SIZE) + { + libpq_append_conn_error(conn, + "compressed protocol message is too large"); + handleFatalError(conn); + return -2; + } + if (conn->compression_buffers_initialized && + conn->compression_buffer.len > 0 && id != PqMsg_CompressedData) + { + libpq_append_conn_error(conn, + "compressed protocol message was interrupted"); + handleFatalError(conn); + return -2; + } +#endif avail = conn->inEnd - conn->inCursor; if (avail < msgLength - 4) { @@ -1887,6 +2395,25 @@ getCopyDataMessage(PGconn *conn) return 0; } +#ifdef USE_ZSTD + if (id == PqMsg_CompressedData) + { + if (!conn->compression_ready) + { + libpq_append_conn_error(conn, + "received compressed data without negotiated compression"); + handleFatalError(conn); + return -2; + } + if (pqDecompressData(conn, msgLength - 4)) + { + handleFatalError(conn); + return -2; + } + continue; + } +#endif + /* * If it's a legitimate async message type, process it. (NOTIFY * messages are not currently possible here, but we handle them for @@ -2135,6 +2662,12 @@ pqEndcopy3(PGconn *conn) if (conn->asyncStatus == PGASYNC_COPY_IN || conn->asyncStatus == PGASYNC_COPY_BOTH) { +#ifdef USE_ZSTD + if (conn->compression_ready && + conn->asyncStatus == PGASYNC_COPY_IN && + pqEndCompressedCopyData(conn) < 0) + return 1; +#endif if (pqPutMsgStart(PqMsg_CopyDone, conn) < 0 || pqPutMsgEnd(conn) < 0) return 1; @@ -2539,6 +3072,12 @@ 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, "zstd") == 0 || + strcmp(conn->compression, "prefer") == 0)) + ADD_STARTUP_OPTION("_pq_.compression", "zstd"); +#endif /* * Add the test_protocol_negotiation option when greasing, to test that diff --git a/src/interfaces/libpq/fe-trace.c b/src/interfaces/libpq/fe-trace.c index 2901fa5b451..93f2fdd4a45 100644 --- a/src/interfaces/libpq/fe-trace.c +++ b/src/interfaces/libpq/fe-trace.c @@ -692,6 +692,11 @@ pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer) pqTraceOutput_CopyData(conn->Pfdebug, message, &logCursor, length, regress); break; + case PqMsg_CompressedData: + fprintf(conn->Pfdebug, "CompressedData"); + /* The compressed payload is intentionally opaque to PQtrace. */ + logCursor = length + 1; + break; case PqMsg_Describe: /* Describe(F) and DataRow(B) use the same identifier. */ Assert(PqMsg_Describe == PqMsg_DataRow); diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h index a737d1db457..fbc005ecee5 100644 --- a/src/interfaces/libpq/libpq-int.h +++ b/src/interfaces/libpq/libpq-int.h @@ -424,6 +424,7 @@ struct pg_conn char *gssdelegation; /* Try to delegate GSS credentials? (0 or 1) */ char *min_protocol_version; /* minimum used protocol version */ char *max_protocol_version; /* maximum used protocol version */ + char *compression; /* protocol compression method */ char *ssl_min_protocol_version; /* minimum TLS protocol version */ char *ssl_max_protocol_version; /* maximum TLS protocol version */ char *target_session_attrs; /* desired session properties */ @@ -579,6 +580,18 @@ struct pg_conn int inStart; /* offset to first unconsumed data in buffer */ int inCursor; /* next byte to tentatively consume */ int inEnd; /* offset to first position after avail data */ +#ifdef USE_ZSTD + void *compression_dctx; /* experimental protocol decompressor */ + void *compression_cctx; /* experimental protocol compressor */ + PQExpBufferData compression_buffer; /* reusable decompression buffer */ + PQExpBufferData compression_output_buffer; + bool compression_buffers_initialized; + bool compression_in_frame; + bool compression_frame_ended; + bool compression_ready; + bool compression_copy_started; +#endif + bool compression_rejected; /* Buffer for data not yet sent to backend */ char *outBuffer; /* currently allocated buffer */ @@ -771,6 +784,9 @@ extern PGresult *PQnfn(PGconn *conn, int fnid, int *result_buf, int buf_size, extern char *pqBuildStartupPacket3(PGconn *conn, int *packetlen, const PQEnvironmentOption *options); extern void pqParseInput3(PGconn *conn); +#ifdef USE_ZSTD +extern void pqCompressionReset(PGconn *conn); +#endif extern int pqGetErrorNotice3(PGconn *conn, bool isError); extern void pqBuildErrorMessage3(PQExpBuffer msg, const PGresult *res, PGVerbosity verbosity, PGContextVisibility show_context); @@ -811,6 +827,11 @@ extern int pqGetInt(int *result, size_t bytes, PGconn *conn); extern int pqPutInt(int value, size_t bytes, PGconn *conn); extern int pqPutMsgStart(char msg_type, PGconn *conn); extern int pqPutMsgEnd(PGconn *conn); +#ifdef USE_ZSTD +extern int pqPutCompressedCopyData(PGconn *conn, const char *buffer, + int nbytes); +extern int pqEndCompressedCopyData(PGconn *conn); +#endif extern int pqReadData(PGconn *conn); extern int pqFlush(PGconn *conn); extern int pqWait(int forRead, int forWrite, PGconn *conn); diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build index b0ae72167a1..a379808b2c8 100644 --- a/src/interfaces/libpq/meson.build +++ b/src/interfaces/libpq/meson.build @@ -161,6 +161,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/t/007_compression.pl b/src/interfaces/libpq/t/007_compression.pl new file mode 100644 index 00000000000..f3b23d5ef95 --- /dev/null +++ b/src/interfaces/libpq/t/007_compression.pl @@ -0,0 +1,182 @@ +# Copyright (c) 2026, PostgreSQL Global Development Group + +use strict; +use warnings FATAL => 'all'; + +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'); + +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_command( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=off', @commands ]); +my ($compressed_stdout, $compressed_stderr) = run_command( + [ '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_command( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=off', + @error_commands ]); +my ($compressed_error_stdout, $compressed_error_stderr) = run_command( + [ '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_command( + [ '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 $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 spans compressed segments and is traced'); + +$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 $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_command( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=invalid', + '-c', 'SELECT 1' ]); +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; +my ($preferred_stdout, $preferred_stderr) = run_command( + [ '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_command( + [ 'psql', '-XAt', '--dbname', $node->connstr('postgres') . ' compression=zstd', + '-c', 'SELECT 1' ]); +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..00c7fd96989 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 + +#include "libpq-int.h" +#include "port/pg_bswap.h" +#endif + static void print_ssl_library(void) { @@ -23,6 +32,479 @@ 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->compression_buffers_initialized || + conn->compression_buffer.len != 0) + { + fprintf(stderr, "server response was not compressed\n"); + PQfinish(conn); + return 1; + } + + /* + * Model an unfinished compressed DataRow without relying on socket + * timing. + */ + appendPQExpBufferChar(&conn->compression_buffer, PqMsg_DataRow); + if (PQExpBufferBroken(&conn->compression_buffer)) + { + fprintf(stderr, "out of memory\n"); + PQfinish(conn); + return 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_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 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; + PQsetTraceFlags(conn, PQTRACE_SUPPRESS_TIMESTAMPS | PQTRACE_REGRESS_MODE); + PQtrace(conn, trace); + + 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 octet_length(data) FROM compression_copy"); + if (PQresultStatus(res) == PGRES_TUPLES_OK && + !strcmp(PQgetvalue(res, 0, 0), "17825792")) + status = 0; + PQclear(res); + +fail: + if (trace != NULL) + { + PQuntrace(conn); + fclose(trace); + } + if (status != 0) + fprintf(stderr, "%s", PQerrorMessage(conn)); + PQfinish(conn); + return status; +} + int main(int argc, char *argv[]) { @@ -31,7 +513,23 @@ 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]); +#ifdef USE_ZSTD + 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-large-datarow CONNINFO, " + "--compression-frame-boundaries CONNINFO, and " + "--compression-pipeline CONNINFO\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, } -- That's all, folks. May the source be with you.