From dcce17bd6eb57cae88f9c0ff9649fa2c6683cb6c Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Mon, 24 Aug 2026 09:01:40 +0200 Subject: Add Frontend support for CompressedMessages Update libpq to handle CompressedMessages. When a CompressedMessages is received, the content is decompressed in connection's decompressBuffer. Functions that are consuming messages (pqParseInput3, getCopyDataMessage and pqFunctionCall3) are updated to process full messages in the decompressBuffer first. --- meson.build | 2 + src/Makefile.global.in | 2 + src/interfaces/libpq/Makefile | 17 +- src/interfaces/libpq/fe-compress-lz4.c | 125 ++++++++++++++ src/interfaces/libpq/fe-compress-zstd.c | 132 ++++++++++++++ src/interfaces/libpq/fe-compress.c | 221 ++++++++++++++++++++++++ src/interfaces/libpq/fe-connect.c | 11 ++ src/interfaces/libpq/fe-protocol3.c | 163 +++++++++++++++-- src/interfaces/libpq/libpq-int.h | 32 ++++ src/interfaces/libpq/meson.build | 3 + 10 files changed, 695 insertions(+), 13 deletions(-) create mode 100644 src/interfaces/libpq/fe-compress-lz4.c create mode 100644 src/interfaces/libpq/fe-compress-zstd.c create mode 100644 src/interfaces/libpq/fe-compress.c diff --git a/meson.build b/meson.build index 67fdf50df02..bd35d1aa372 100644 --- a/meson.build +++ b/meson.build @@ -3630,7 +3630,9 @@ libpq_deps += [ gssapi, ldap_r, libintl, + lz4, ssl, + zstd, ] libpq_oauth_deps += [ diff --git a/src/Makefile.global.in b/src/Makefile.global.in index 46f88de6f50..499ebf3c0a6 100644 --- a/src/Makefile.global.in +++ b/src/Makefile.global.in @@ -201,9 +201,11 @@ with_liburing = @with_liburing@ with_libxml = @with_libxml@ with_libxslt = @with_libxslt@ with_llvm = @with_llvm@ +with_lz4 = @with_lz4@ with_system_tzdata = @with_system_tzdata@ with_uuid = @with_uuid@ with_zlib = @with_zlib@ +with_zstd = @with_zstd@ enable_rpath = @enable_rpath@ enable_nls = @enable_nls@ enable_debug = @enable_debug@ diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 0963995eed4..bbfb67aacca 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -15,7 +15,7 @@ subdir = src/interfaces/libpq top_builddir = ../../.. include $(top_builddir)/src/Makefile.global -export with_ssl with_gssapi with_krb_srvnam +export with_ssl with_gssapi with_krb_srvnam with_lz4 with_zstd PGFILEDESC = "PostgreSQL Access Library" @@ -33,6 +33,9 @@ OBJS = \ $(WIN32RES) \ fe-auth-scram.o \ fe-cancel.o \ + fe-compress.o \ + fe-compress-zstd.o \ + fe-compress-lz4.o \ fe-connect.o \ fe-exec.o \ fe-lobj.o \ @@ -87,9 +90,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 -llz4 -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 -llz4 -lzstd $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE) endif ifeq ($(PORTNAME), win32) SHLIB_LINK += -lshell32 -lws2_32 -lsecur32 $(filter -lcomerr32 -lkrb5_32, $(LIBS)) @@ -109,6 +112,14 @@ $(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,libssl) $(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,libcrypto) endif +ifeq ($(with_lz4),yes) +$(call add_to_list,PKG_CONFIG_REQUIRES_PRIVATE,liblz4) +endif + +ifeq ($(with_zstd),yes) +$(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-compress-lz4.c b/src/interfaces/libpq/fe-compress-lz4.c new file mode 100644 index 00000000000..a78838127a1 --- /dev/null +++ b/src/interfaces/libpq/fe-compress-lz4.c @@ -0,0 +1,125 @@ +/*------------------------------------------------------------------------- + * + * fe-compress-lz4.c + * Lz4 decompression support for frontend/backend protocol + * + * Portions Copyright (c) 2026, PostgreSQL Global Development Group + * + * + * IDENTIFICATION + * src/interfaces/libpq/fe-compress-lz4.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include + +#include "libpq-fe.h" +#include "libpq-int.h" + +#ifndef USE_LZ4 + +int +pqInitDecompressorLz4(PGconn *conn) +{ + libpq_append_conn_error(conn, "client does not support compression with lz4"); + return EOF; +} + +#else + +#include + +#define LZ4_CHUNK_SZ 64 * 1024 /* 64kB as maximum chunk size read */ + +static void free_context_lz4(PGconn *conn); +static int decompress_payload_lz4(PGconn *conn); + +static const pqDecompressor pqDecompressorLz4 = { + .decompress_payload = decompress_payload_lz4, + .free_context = free_context_lz4, +}; + +static void +free_context_lz4(PGconn *conn) +{ + LZ4F_decompressionContext_t dctx; + + Assert(conn->compress_state); + dctx = (LZ4F_decompressionContext_t) conn->compress_state; + LZ4F_freeDecompressionContext(dctx); + conn->compress_state = NULL; +} + +/* + * Initialize lz4 decompression context + * + * returns 0 if OK, EOF if trouble + */ +int +pqInitDecompressorLz4(PGconn *conn) +{ + LZ4F_decompressionContext_t dctx; + LZ4F_errorCode_t ctxError; + + ctxError = LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION); + if (LZ4F_isError(ctxError)) + { + libpq_append_conn_error(conn, "out of memory"); + return EOF; + } + + conn->compress_state = dctx; + conn->decompressor = pqDecompressorLz4; + conn->decompress_chunk_size = LZ4_CHUNK_SZ; + return 0; +} + +/* + * Decompress the compressed payload using lz4. The result will be + * stored in conn->decompressBuffer. + * + * returns 0 if OK, EOF if trouble + */ +static int +decompress_payload_lz4(PGconn *conn) +{ + size_t res; + msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *outBuf = &conn->decompressBuffer; + LZ4F_decompressionContext_t dctx = (LZ4F_decompressionContext_t) conn->compress_state; + + do + { + void *srcBuffer = msgBuf->buffer + msgBuf->start + conn->compress_cursor; + void *dstBuffer = outBuf->buffer + outBuf->end; + size_t consumed = conn->compress_end - conn->compress_cursor; + size_t decompressed = outBuf->bufSize - outBuf->end; + + Assert(msgBuf->start + conn->compress_end <= msgBuf->bufSize); + res = LZ4F_decompress(dctx, dstBuffer, &decompressed, + srcBuffer, &consumed, NULL); + + if (LZ4F_isError(res)) + { + libpq_append_conn_error(conn, "could not decompress data: %s", + LZ4F_getErrorName(res)); + return EOF; + } + + /* Update msgBuf with the consumed bytes */ + Assert(msgBuf->start + conn->compress_end + consumed <= msgBuf->bufSize); + conn->compress_cursor += consumed; + + /* And update decompressBuffer with the decompressed bytes */ + Assert(outBuf->end + decompressed <= outBuf->bufSize); + outBuf->end += decompressed; + if (outBuf->end == outBuf->bufSize) + /* Output buffer is full, leave the message unconsumed */ + return 0; + } while (conn->compress_cursor < conn->compress_end); + return 0; +} + +#endif diff --git a/src/interfaces/libpq/fe-compress-zstd.c b/src/interfaces/libpq/fe-compress-zstd.c new file mode 100644 index 00000000000..9e99e7b7c5c --- /dev/null +++ b/src/interfaces/libpq/fe-compress-zstd.c @@ -0,0 +1,132 @@ +/*------------------------------------------------------------------------- + * + * fe-compress-zstd.c + * Zstd decompression support for frontend/backend protocol + * + * Portions Copyright (c) 2026, PostgreSQL Global Development Group + * + * + * IDENTIFICATION + * src/interfaces/libpq/fe-compress-zstd.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include "libpq-fe.h" +#include "libpq-int.h" + +#ifndef USE_ZSTD + +int +pqInitDecompressorZstd(PGconn *conn) +{ + libpq_append_conn_error(conn, "client does not support compression with zstd"); + return EOF; +} + +#else + +#include +#include + +static void free_context_zstd(PGconn *conn); +static int decompress_payload_zstd(PGconn *conn); + +static const pqDecompressor pqDecompressorZstd = { + .decompress_payload = decompress_payload_zstd, + .free_context = free_context_zstd, +}; + +static void +free_context_zstd(PGconn *conn) +{ + ZSTD_DStream *dctx; + + Assert(conn->compress_state); + dctx = (ZSTD_DStream *) conn->compress_state; + ZSTD_freeDStream(dctx); + conn->compress_state = NULL; +} + +/* + * Initialize zstd decompression context + * + * returns 0 if OK, EOF if trouble + */ +int +pqInitDecompressorZstd(PGconn *conn) +{ + ZSTD_DStream *dctx; + + dctx = ZSTD_createDStream(); + + if (dctx == NULL) + { + libpq_append_conn_error(conn, "out of memory"); + return EOF; + } + conn->compress_state = dctx; + conn->decompressor = pqDecompressorZstd; + conn->decompress_chunk_size = ZSTD_DStreamOutSize(); + return 0; +} + +/* + * Decompress the compressed payload using zstd. The result will be + * stored in conn->decompressBuffer. + * + * returns 0 if OK, EOF if trouble + */ +static int +decompress_payload_zstd(PGconn *conn) +{ + size_t res; + ZSTD_inBuffer inBuf = {conn->inBuffer.buffer + conn->inBuffer.start, conn->compress_end, conn->compress_cursor}; + ZSTD_outBuffer outBuf; + ZSTD_DStream *dctx = (ZSTD_DStream *) conn->compress_state; + + outBuf.dst = conn->decompressBuffer.buffer; + outBuf.size = conn->decompressBuffer.bufSize; + outBuf.pos = conn->decompressBuffer.end; + + /* + * It's possible to have inBuf.pos == inBuf.size and outBuf.pos == + * outBuf.size after a ZSTD_decompressStream due to how frames are closed. + * zstd normally keeps the last byte "hostage" when the last block of a + * frame is decompressed, forcing (inBuf.pos < inBuf.size) to be true. + * + * However, as we leave the frame opened, this never happens, and we can + * reach a point where the block is completely consumed while the outBuf + * is full. + * + * Thus the use of the do while as we want to call ZSTD_decompressStream + * even if inBuf was completely processed, to force zstd to flush any + * leftover buffers. + */ + do + { + res = ZSTD_decompressStream(dctx, &outBuf, &inBuf); + + if (ZSTD_isError(res)) + { + libpq_append_conn_error(conn, "could not decompress data: %s", + ZSTD_getErrorName(res)); + return EOF; + } + + if (outBuf.pos == outBuf.size) + + /* + * output buffer is full, break to let the outer loop either + * consume the messages or enlarge the buffer + */ + break; + } while (inBuf.pos < inBuf.size); + + conn->decompressBuffer.end = outBuf.pos; + conn->compress_cursor = inBuf.pos; + return 0; +} + +#endif diff --git a/src/interfaces/libpq/fe-compress.c b/src/interfaces/libpq/fe-compress.c new file mode 100644 index 00000000000..eff1cb52098 --- /dev/null +++ b/src/interfaces/libpq/fe-compress.c @@ -0,0 +1,221 @@ +/*------------------------------------------------------------------------- + * + * fe-compress.c + * Decompression support for frontend/backend protocol + * + * Portions Copyright (c) 2026, PostgreSQL Global Development Group + * + * + * IDENTIFICATION + * src/interfaces/libpq/fe-compress.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include + +#include "common/compression.h" +#include "libpq-fe.h" +#include "libpq-int.h" + +static int pqGetMsgLen(PGconn *conn, msg_buffer * msgBuf); +static bool pqHasFullMessage(PGconn *conn, msg_buffer * msgBuf); + +/* + * pqReadCompressedMessage: Read the content of a CompressedMessages message. + * + * returns EOF if not enough data, -2 if trouble, 0 if CompressedMessages + * content has been fully processed, 1 if CompressedMessages has been + * successfully decompressed + */ +int +pqReadCompressedMessage(PGconn *conn, int msgLength) +{ + char compress_algorithm; + int endCompressMessage = conn->inBuffer.start + msgLength + 5; + + if (conn->compress_cursor > 0 + && conn->compress_cursor == msgLength + 5) + { + /* + * The CompressedMessages was fully processed. Advance inBuffer so the + * message can be tagged as consumed. + */ + conn->inBuffer.cursor = endCompressMessage; + conn->compress_cursor = 0; + return 0; + } + + /* Get compression algorithm */ + if (pqGetc(&compress_algorithm, conn, &conn->inBuffer) == EOF) + return EOF; + + if (compress_algorithm != PG_COMPRESSION_ZSTD && + compress_algorithm != PG_COMPRESSION_LZ4) + { + libpq_append_conn_error( + conn, "invalid compression algorithm in CompressedMessage: %d", + compress_algorithm); + return 0; + } + + /* get the message types */ + if (pqGets(&conn->workBuffer, conn, &conn->inBuffer) == EOF) + return EOF; + + if (conn->compress_algorithm != compress_algorithm) + { + /* + * Compression in the message is different from our current + * compression context, free the previous context to start from a + * clean slate + */ + if (conn->decompressor.free_context != NULL) + conn->decompressor.free_context(conn); + switch (compress_algorithm) + { + case PG_COMPRESSION_LZ4: + if (pqInitDecompressorLz4(conn)) + return -2; + break; + case PG_COMPRESSION_ZSTD: + if (pqInitDecompressorZstd(conn)) + return -2; + break; + case PG_COMPRESSION_GZIP: + case PG_COMPRESSION_NONE: + pg_unreachable(); + } + conn->compress_algorithm = compress_algorithm; + } + + conn->compress_cursor = conn->inBuffer.cursor - conn->inBuffer.start; + conn->compress_end = endCompressMessage - conn->inBuffer.start; + return 1; +} + +/* + * pqDecompressPayload: Decompress into the decompress buffer until: + * - A full message is available for processing. + * Or + * - All input bytes are processed and the decompress buffer isn't full. + * + * Due to keeping the compression frame opened, it's possible to have + * cursor==end while the compressor still has bytes to output, but couldn't due + * to decompressBuffer being full. We rely on decompressBuffer being full as a + * condition to run one more pq_decompress_payload to make sure there's nothing + * stuck in the decompressor's buffers. + * + * returns EOF if not enough data, -2 if trouble, 0 if there's at least one + * full message available + */ +int +pqDecompressPayload(PGconn *conn) +{ + bool outBufferFull; + + if (pqHasFullMessage(conn, &conn->decompressBuffer)) + { + /* There's already a full message in decompress buffer */ + return 0; + } + + outBufferFull = conn->decompressBuffer.end == conn->decompressBuffer.bufSize; + while (conn->compress_cursor < conn->compress_end || outBufferFull) + { + if (outBufferFull) + { + int msgLen = pqGetMsgLen(conn, &conn->decompressBuffer); + size_t bytes_needed; + + if (msgLen == -1) + { + /* + * Not enough bytes for the msg length. Either the last + * message is in the last 5 bytes of decompressBuffer, or + * everything was processed. Just ask for 4 bytes, and + * pqCheckMsgBufferSpace will clean the already processed + * messages. + */ + msgLen = 4; + } + /* bytes_needed needs to include the starting 1 byte id. */ + bytes_needed = Max(msgLen + 1, conn->decompress_chunk_size); + + /* + * decompressBuffer is full, and we don't have a full message, + * call pqCheckMsgBufferSpace to either left justify the content + * or enlarge the buffer. Once space is freed, decompress needs to + * be called again to give the decompressor a chance to flush its + * internal buffers. + */ + if (pqCheckMsgBufferSpace(bytes_needed + conn->decompressBuffer.start, &conn->decompressBuffer, conn)) + { + return -2; + } + } + + /* + * We should always have available bytes before calling + * decompress_payload + */ + Assert(conn->decompressBuffer.bufSize > conn->decompressBuffer.end); + + /* Do the decompression */ + if (conn->decompressor.decompress_payload(conn)) + + /* + * there was a fatal error while decompressing the payload, bail + * out + */ + return -2; + + if (pqHasFullMessage(conn, &conn->decompressBuffer)) + + /* + * We have at least one full message, exit to let the outer loop + * consume it + */ + return 0; + + outBufferFull = conn->decompressBuffer.end == conn->decompressBuffer.bufSize; + } + return EOF; +} + +/* + * Returns the length of the next message in the msgBuf + * + * Returns EOF if there's not enough bytes to read the length + */ +static int +pqGetMsgLen(PGconn *conn, msg_buffer * msgBuf) +{ + char id; + int msgLength; + + msgBuf->cursor = msgBuf->start; + if (pqGetc(&id, conn, msgBuf)) + return EOF; + if (pqGetInt(&msgLength, 4, conn, msgBuf)) + return EOF; + + return msgLength; +} + +/* + * Returns true if there's a full message in the msgBuf + */ +static bool +pqHasFullMessage(PGconn *conn, msg_buffer * msgBuf) +{ + int len = pqGetMsgLen(conn, msgBuf); + int available = msgBuf->end - msgBuf->start; + + if (len < 0) + return false; + + /* Available bytes needs to include the 1 byte id */ + return available >= 1 + len; +} diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c index 829a68be557..599b917d4fe 100644 --- a/src/interfaces/libpq/fe-connect.c +++ b/src/interfaces/libpq/fe-connect.c @@ -554,7 +554,11 @@ pqDropConnection(PGconn *conn, bool flushInput) /* Optionally discard any unread data */ if (flushInput) + { conn->inBuffer.start = conn->inBuffer.cursor = conn->inBuffer.end = 0; + if (conn->decompressBuffer.buffer) + conn->decompressBuffer.start = conn->decompressBuffer.cursor = conn->decompressBuffer.end = 0; + } /* Always discard any unsent data */ conn->outCount = 0; @@ -2793,6 +2797,10 @@ pqConnectDBStart(PGconn *conn) /* Ensure our buffers are empty */ conn->inBuffer.start = conn->inBuffer.cursor = conn->inBuffer.end = 0; + if (conn->decompressBuffer.buffer) + conn->decompressBuffer.start = conn->decompressBuffer.cursor = conn->decompressBuffer.end = 0; + if (conn->inBuffer.buffer) + conn->inBuffer.start = conn->inBuffer.cursor = conn->inBuffer.end = 0; conn->outCount = 0; /* @@ -5218,6 +5226,9 @@ freePGconn(PGconn *conn) free(conn->inBuffer.buffer); free(conn->outBuffer); free(conn->rowBuf); + if (conn->decompressor.free_context) + conn->decompressor.free_context(conn); + free(conn->decompressBuffer.buffer); termPQExpBuffer(&conn->errorMessage); termPQExpBuffer(&conn->workBuffer); diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c index b34e14d5b81..fbb563d850d 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_CompressedMessages || \ (id) == PqMsg_ParameterDescription) @@ -61,6 +62,45 @@ static void reportErrorPosition(PQExpBuffer msg, const char *query, static size_t build_startup_packet(const PGconn *conn, char *packet, const PQEnvironmentOption *options); +/* + * getMsgBuffer: returns the msg_buffer to consume. + * + * If compression is used, pqDecompressPayload is called to process any + * available bytes. If we have a full decompressed message after this, + * decompressBuffer will be returned as the buffer to consume. + * Otherwise, we return inBuffer. + * + * returns the buffer to use, or NULL on error. + */ +static msg_buffer * +getMsgBuffer(PGconn *conn) +{ + int ret; + + if (conn->compress_algorithm == PG_COMPRESSION_NONE + || conn->compress_cursor == 0) + { + /* No compression, just use inBuffer */ + return &conn->inBuffer; + } + + /* Compression is enabled, try to decompress */ + ret = pqDecompressPayload(conn); + if (ret == 0) + + /* + * There's at least one full message available in the decompression + * buffer, process it + */ + return &conn->decompressBuffer; + else if (ret == -2) + { + handleFatalError(conn); + return NULL; + } + + return &conn->inBuffer; +} /* * parseInput: if appropriate, parse input data from backend @@ -73,13 +113,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. @@ -105,6 +149,18 @@ pqParseInput3(PGconn *conn) handleSyncLoss(conn, id, msgLength); return; } + if (id != PqMsg_CompressedMessages + && conn->decompressBuffer.start != conn->decompressBuffer.end + && msgBuf != &conn->decompressBuffer) + { + /* + * If we have data in the decompress buffer, only a + * PqMsg_CompressedMessages is acceptable. Anything else is a + * protocol issue + */ + handleSyncLoss(conn, id, msgLength); + return; + } /* * Can't process if message body isn't all here yet. @@ -139,6 +195,11 @@ pqParseInput3(PGconn *conn) * NOTIFY and NOTICE messages can happen in any state; always process * them right away. * + * CompressedMessages is a special case. As the nested messages are + * processed first, it's possible to have a CompressedMessages while + * in IDLE state during the second pass. Process them if it's only to + * tag the message as consumed. + * * Most other messages should only be processed while in BUSY state. * (In particular, in READY state we hold off further parsing until * the application collects the current PGresult.) @@ -161,6 +222,17 @@ pqParseInput3(PGconn *conn) if (pqGetErrorNotice3(conn, false, msgBuf)) return; } + else if (id == PqMsg_CompressedMessages + && conn->compress_cursor > 0 + && conn->compress_cursor == conn->inBuffer.start + msgLength + 5) + { + /* + * The CompressedMessages was fully processed. Advance inBuffer so + * the message can be tagged as consumed. + */ + conn->inBuffer.cursor = conn->inBuffer.start + msgLength + 5; + conn->compress_cursor = 0; + } else if (conn->asyncStatus != PGASYNC_BUSY) { /* If not IDLE state, just wait ... */ @@ -285,6 +357,24 @@ pqParseInput3(PGconn *conn) case PqMsg_BindComplete: /* Nothing to do for this message type */ break; + case PqMsg_CompressedMessages: + { + int ret = pqReadCompressedMessage(conn, msgLength); + + if (ret == EOF) + return; + else if (ret == -2) + handleFatalError(conn); + else if (ret == 1) + + /* + * CompressedMessages is intentionally not + * consumed as the decompress content needs to be + * processed first + */ + continue; + break; + } case PqMsg_CloseComplete: /* If we're doing PQsendClose, we're done; else ignore */ if (conn->cmd_queue_head && @@ -1243,9 +1333,9 @@ reportErrorPosition(PQExpBuffer msg, const char *query, int loc, int encoding) * * The only caller of reportErrorPosition() is pqBuildErrorMessage3(); it * gets its query from either a PQresultErrorField() or a PGcmdQueueEntry, - * both of which must have fit into inBuffer/outBuffer. So slen fits - * inside an int, but we can't assume that (slen * sizeof(int)) fits - * inside a size_t. + * both of which must have fit into + * conn->inBuffer/decompressBuffer/outBuffer. So slen fits inside an int, + * but we can't assume that (slen * sizeof(int)) fits inside a size_t. */ slen = strlen(wquery) + 1; if (slen > SIZE_MAX / sizeof(int)) @@ -1811,12 +1901,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 @@ -1832,6 +1925,13 @@ getCopyDataMessage(PGconn *conn, msg_buffer * *out_buf) handleSyncLoss(conn, id, msgLength); return -2; } + if (id != PqMsg_CompressedMessages + && conn->decompressBuffer.start != conn->decompressBuffer.end + && msgBuf != &conn->decompressBuffer) + { + handleSyncLoss(conn, id, msgLength); + return -2; + } avail = msgBuf->end - msgBuf->cursor; if (avail < msgLength - 4) { @@ -1866,6 +1966,24 @@ getCopyDataMessage(PGconn *conn, msg_buffer * *out_buf) if (getNotify(conn, msgBuf)) return 0; break; + case PqMsg_CompressedMessages: + { + int ret = pqReadCompressedMessage(conn, msgLength - 4); + + if (ret == EOF) + return 0; + else if (ret == -2) + handleFatalError(conn); + else if (ret == 1) + + /* + * CompressedMessages is intentionally not consumed as + * the decompress content needs to be processed first + */ + continue; + + break; + } case PqMsg_NoticeResponse: if (pqGetErrorNotice3(conn, false, msgBuf)) return 0; @@ -1924,7 +2042,7 @@ int pqGetCopyData3(PGconn *conn, char **buffer, int async) { int msgLength; - msg_buffer *msgBuf = &conn->inBuffer; + msg_buffer *msgBuf; for (;;) { @@ -2036,7 +2154,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) @@ -2192,7 +2310,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 +2368,9 @@ pqFunctionCall3(PGconn *conn, Oid fnid, */ needInput = true; + msgBuf = getMsgBuffer(conn); + if (!msgBuf) + break; msgBuf->cursor = msgBuf->start; if (pqGetc(&id, conn, msgBuf)) continue; @@ -2271,6 +2392,13 @@ pqFunctionCall3(PGconn *conn, Oid fnid, handleSyncLoss(conn, id, msgLength); break; } + if (id != PqMsg_CompressedMessages + && conn->decompressBuffer.start != conn->decompressBuffer.end + && msgBuf != &conn->decompressBuffer) + { + handleSyncLoss(conn, id, msgLength); + break; + } /* * Can't process if message body isn't all here yet. @@ -2342,6 +2470,21 @@ pqFunctionCall3(PGconn *conn, Oid fnid, continue; status = PGRES_FATAL_ERROR; break; + case PqMsg_CompressedMessages: + { + int ret = pqReadCompressedMessage(conn, msgLength); + + if (ret == EOF) + continue; + else if (ret == -2) + handleFatalError(conn); + else if (ret == 1) + { + needInput = false; + continue; + } + break; + } case PqMsg_NotificationResponse: /* handle notify and go back to processing return values */ if (getNotify(conn, msgBuf)) diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h index 06e8ff68557..6e6d78d0f46 100644 --- a/src/interfaces/libpq/libpq-int.h +++ b/src/interfaces/libpq/libpq-int.h @@ -40,6 +40,7 @@ /* include stuff common to fe and be */ #include "libpq/pqcomm.h" +#include "common/compression.h" /* include stuff found in fe only */ #include "fe-auth-sasl.h" #include "pqexpbuffer.h" @@ -375,6 +376,12 @@ typedef struct msg_buffer int end; /* offset to first position after avail data */ } msg_buffer; +typedef struct pqDecompressor +{ + int (*decompress_payload) (PGconn *conn); + void (*free_context) (PGconn *conn); +} pqDecompressor; + /* * PGconn stores all the state data associated with a single connection * to a backend. @@ -588,6 +595,18 @@ struct pg_conn * code that uses ints during size calculations. */ msg_buffer inBuffer; + int compress_cursor; /* Location of currently processed + * compressed bytes, relative to + * inBuffer's start */ + int compress_end; /* Location of end of CompressedMessages, + * relative to inBuffer's start */ + + msg_buffer decompressBuffer; /* Buffer for the decompressed messages */ + + pg_compress_algorithm compress_algorithm; /* Compression used */ + pqDecompressor decompressor; + int decompress_chunk_size; + void *compress_state; /* private state for compression */ /* Buffer for data not yet sent to backend */ @@ -799,6 +818,19 @@ extern PGresult *pqFunctionCall3(PGconn *conn, Oid fnid, extern int PQsendCancelRequest(PGconn *cancelConn); +/* === in fe-compress.c === */ + +extern int pqReadCompressedMessage(PGconn *conn, int msgLength); +extern int pqDecompressPayload(PGconn *conn); + +/* === in fe-compress-zstd.c === */ + +extern int pqInitDecompressorZstd(PGconn *conn); + +/* === in fe-compress-lz4.c === */ + +extern int pqInitDecompressorLz4(PGconn *conn); + /* === in fe-misc.c === */ /* diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build index b0ae72167a1..83a085ec9b1 100644 --- a/src/interfaces/libpq/meson.build +++ b/src/interfaces/libpq/meson.build @@ -5,6 +5,9 @@ libpq_sources = files( 'fe-auth-scram.c', 'fe-auth.c', 'fe-cancel.c', + 'fe-compress.c', + 'fe-compress-zstd.c', + 'fe-compress-lz4.c', 'fe-connect.c', 'fe-exec.c', 'fe-lobj.c', -- 2.50.1 (Apple Git-155)