From 7afbca00d52fd8512e53eede651bf9ebecb25638 Mon Sep 17 00:00:00 2001 From: "Chao Li (Evan)" Date: Fri, 14 Aug 2026 11:29:21 +0800 Subject: [PATCH v9 2/2] Fix detection of truncated zstd and LZ4 dump data `pg_restore` did not verify that zstd and LZ4 decompression completed the final frame when reading custom and directory-format dumps. As a result, it could accept a dump whose compressed data was truncated. Track final-frame completion and reject incomplete compressed input. For zstd, flush any internally buffered output before performing the final check. Add regression tests for truncated zstd and LZ4 data in custom and directory-format dumps. Author: Chao Li Co-authored-by: Daniel Gustafsson Reviewed-by: Daniel Gustafsson Discussion: https://postgr.es/m/5962B878-C43D-4EBC-9E95-1F945CE5E586@gmail.com --- src/bin/pg_dump/compress_gzip.c | 6 +- src/bin/pg_dump/compress_lz4.c | 50 ++++++--- src/bin/pg_dump/compress_zstd.c | 67 +++++++++++- src/bin/pg_dump/t/006_pg_dump_compress.pl | 125 ++++++++++++++++++++++ 4 files changed, 228 insertions(+), 20 deletions(-) diff --git a/src/bin/pg_dump/compress_gzip.c b/src/bin/pg_dump/compress_gzip.c index bd1043b2366..6781c41ec2f 100644 --- a/src/bin/pg_dump/compress_gzip.c +++ b/src/bin/pg_dump/compress_gzip.c @@ -205,7 +205,8 @@ ReadDataFromArchiveGzip(ArchiveHandle *AH, CompressorState *cs) res = inflate(zp, 0); if (res != Z_OK && res != Z_STREAM_END) - pg_fatal("could not uncompress data: %s", zp->msg); + pg_fatal("could not uncompress data: %s", + zp->msg ? zp->msg : "unknown error"); out[DEFAULT_IO_BUFFER_SIZE - zp->avail_out] = '\0'; ahwrite(out, 1, DEFAULT_IO_BUFFER_SIZE - zp->avail_out, AH); @@ -220,7 +221,8 @@ ReadDataFromArchiveGzip(ArchiveHandle *AH, CompressorState *cs) zp->avail_out = DEFAULT_IO_BUFFER_SIZE; res = inflate(zp, 0); if (res != Z_OK && res != Z_STREAM_END) - pg_fatal("could not uncompress data: %s", zp->msg); + pg_fatal("could not uncompress data: %s", + zp->msg ? zp->msg : "unknown error"); out[DEFAULT_IO_BUFFER_SIZE - zp->avail_out] = '\0'; ahwrite(out, 1, DEFAULT_IO_BUFFER_SIZE - zp->avail_out, AH); diff --git a/src/bin/pg_dump/compress_lz4.c b/src/bin/pg_dump/compress_lz4.c index 500d5e16a6d..2dba2eceab8 100644 --- a/src/bin/pg_dump/compress_lz4.c +++ b/src/bin/pg_dump/compress_lz4.c @@ -58,6 +58,7 @@ typedef struct LZ4State * decompression operations. */ bool compressing; + bool frame_finished; /* * I/O buffer area. @@ -160,6 +161,12 @@ ReadDataFromArchiveLZ4(ArchiveHandle *AH, CompressorState *cs) LZ4F_decompressOptions_t dec_opt; LZ4F_errorCode_t status; + /* + * cs->private_data is an LZ4State for compression, whereas this function + * uses a short-lived decompression context. Keep its state local. + */ + bool dec_done = false; + memset(&dec_opt, 0, sizeof(dec_opt)); status = LZ4F_createDecompressionContext(&ctx, LZ4F_VERSION); if (LZ4F_isError(status)) @@ -187,12 +194,16 @@ ReadDataFromArchiveLZ4(ArchiveHandle *AH, CompressorState *cs) if (LZ4F_isError(status)) pg_fatal("could not decompress: %s", LZ4F_getErrorName(status)); + dec_done = (status == 0); ahwrite(outbuf, 1, out_size, AH); readp += read_size; } } + if (!dec_done) + pg_fatal("could not decompress data: compressed stream is incomplete"); + pg_free(outbuf); pg_free(readbuf); @@ -414,11 +425,7 @@ LZ4Stream_read_internal(LZ4State *state, void *ptr, int ptrsize, bool eol_flag) /* Lazy init */ if (!LZ4Stream_init(state, false /* decompressing */ )) - { - pg_log_error("unable to initialize LZ4 library: %s", - LZ4F_getErrorName(state->errcode)); return -1; - } /* Loop until postcondition is satisfied */ while (remaining > 0) @@ -466,12 +473,17 @@ LZ4Stream_read_internal(LZ4State *state, void *ptr, int ptrsize, bool eol_flag) rsize = fread(state->buffer, 1, state->buflen, state->fp); if (rsize < state->buflen && !feof(state->fp)) - { - pg_log_error("could not read from input file: %m"); return -1; - } + if (rsize == 0) + { + if (!state->frame_finished) + { + errno = EIO; + return -1; + } break; /* must be EOF */ + } state->bufdata = rsize; state->bufnext = 0; } @@ -492,10 +504,9 @@ LZ4Stream_read_internal(LZ4State *state, void *ptr, int ptrsize, bool eol_flag) if (LZ4F_isError(status)) { state->errcode = status; - pg_log_error("could not read from input file: %s", - LZ4F_getErrorName(state->errcode)); return -1; } + state->frame_finished = (status == 0); state->bufnext += inlen; state->outbufdata = outlen; state->outbufnext = 0; @@ -600,12 +611,16 @@ LZ4Stream_gets(char *ptr, int size, CompressFileHandle *CFH) ret = LZ4Stream_read_internal(state, ptr, size - 1, true); - /* - * LZ4Stream_read_internal returning 0 or -1 means that it was either an - * EOF or an error, but gets_func is defined to return NULL in either case - * so we can treat both the same here. - */ - if (ret <= 0) + if (ret < 0) + { + /* gets_func must return NULL rather than exiting on an error. */ + pg_log_error("could not read from input file: %s", + LZ4Stream_get_error(CFH)); + return NULL; + } + + /* gets_func returns NULL on EOF when no characters have been read. */ + if (ret == 0) return NULL; /* @@ -681,6 +696,11 @@ LZ4Stream_close(CompressFileHandle *CFH) } else { + if (!state->frame_finished) + { + errno = EIO; + success = false; + } status = LZ4F_freeDecompressionContext(state->dtx); if (LZ4F_isError(status)) { diff --git a/src/bin/pg_dump/compress_zstd.c b/src/bin/pg_dump/compress_zstd.c index 68f1d815917..22e6096c367 100644 --- a/src/bin/pg_dump/compress_zstd.c +++ b/src/bin/pg_dump/compress_zstd.c @@ -45,6 +45,7 @@ typedef struct ZstdCompressorState ZSTD_DStream *dstream; ZSTD_outBuffer output; ZSTD_inBuffer input; + bool frame_finished; /* pointer to a static string like from strerror(), for Zstd_write() */ const char *zstderror; @@ -165,6 +166,7 @@ ReadDataFromArchiveZstd(ArchiveHandle *AH, CompressorState *cs) ZSTD_inBuffer *input = &zstdcs->input; size_t input_allocated_size = ZSTD_DStreamInSize(); size_t res; + bool frame_finished = false; for (;;) { @@ -193,6 +195,7 @@ ReadDataFromArchiveZstd(ArchiveHandle *AH, CompressorState *cs) res = ZSTD_decompressStream(zstdcs->dstream, output, input); if (ZSTD_isError(res)) pg_fatal("could not decompress data: %s", ZSTD_getErrorName(res)); + frame_finished = (res == 0); /* * then write the decompressed data to the output handle @@ -200,10 +203,36 @@ ReadDataFromArchiveZstd(ArchiveHandle *AH, CompressorState *cs) ((char *) output->dst)[output->pos] = '\0'; ahwrite(output->dst, 1, output->pos, AH); - if (res == 0) + if (frame_finished) break; /* End of frame */ } } + + if (!frame_finished) + { + ZSTD_inBuffer empty = {NULL, 0, 0}; + + /* + * A full output buffer with a positive return value might leave data + * in zstd's internal buffers. Call the decompressor with empty input + * until it has flushed that data. + */ + while (!frame_finished) + { + output->pos = 0; + res = ZSTD_decompressStream(zstdcs->dstream, output, &empty); + if (ZSTD_isError(res)) + pg_fatal("could not decompress data: %s", + ZSTD_getErrorName(res)); + frame_finished = (res == 0); + + if (output->pos == 0 && !frame_finished) + pg_fatal("could not decompress data: compressed stream is incomplete"); + + ((char *) output->dst)[output->pos] = '\0'; + ahwrite(output->dst, 1, output->pos, AH); + } + } } /* Public routine that supports Zstd compressed data I/O */ @@ -318,9 +347,33 @@ Zstd_read_internal(void *ptr, size_t size, CompressFileHandle *CFH, bool exit_on Assert(cnt <= input_allocated_size); - /* If we have no more input to consume, we're done */ + /* If we have no more input to consume, verify the final frame. */ if (cnt == 0) + { + ZSTD_inBuffer empty = {NULL, 0, 0}; + + if (!zstdcs->frame_finished) + { + res = ZSTD_decompressStream(zstdcs->dstream, output, &empty); + if (ZSTD_isError(res)) + { + if (exit_on_error) + pg_fatal("could not decompress data: %s", ZSTD_getErrorName(res)); + return -1; + } + + zstdcs->frame_finished = (res == 0); + if (!zstdcs->frame_finished && output->pos == 0) + { + zstdcs->zstderror = "compressed stream is incomplete"; + if (exit_on_error) + pg_fatal("could not decompress data: %s", zstdcs->zstderror); + return -1; + } + } + break; + } } while (input->pos < input->size) @@ -335,10 +388,12 @@ Zstd_read_internal(void *ptr, size_t size, CompressFileHandle *CFH, bool exit_on return -1; } + zstdcs->frame_finished = (res == 0); + if (output->pos == output->size) break; /* No more room for output */ - if (res == 0) + if (zstdcs->frame_finished) break; /* End of frame */ } } @@ -477,6 +532,12 @@ Zstd_close(CompressFileHandle *CFH) if (zstdcs->dstream) { + if (!zstdcs->frame_finished) + { + errno = EIO; + zstdcs->zstderror = "compressed stream is incomplete"; + success = false; + } ZSTD_freeDStream(zstdcs->dstream); pg_free(unconstify(void *, zstdcs->input.src)); } diff --git a/src/bin/pg_dump/t/006_pg_dump_compress.pl b/src/bin/pg_dump/t/006_pg_dump_compress.pl index d4ce6b18077..f59d9bf0539 100644 --- a/src/bin/pg_dump/t/006_pg_dump_compress.pl +++ b/src/bin/pg_dump/t/006_pg_dump_compress.pl @@ -12,6 +12,7 @@ use strict; use warnings FATAL => 'all'; +use File::Copy qw(copy); use PostgreSQL::Test::Cluster; use PostgreSQL::Test::Utils; use Test::More; @@ -38,6 +39,75 @@ my $supports_gzip = check_pg_config("#define HAVE_LIBZ 1"); my $supports_lz4 = check_pg_config("#define USE_LZ4 1"); my $supports_zstd = check_pg_config("#define USE_ZSTD 1"); +sub +truncate_custom_compressed_data +{ + my ($path, $magic, $empty) = @_; + my ($data, $pos, $first_lenpos, $last_lenpos, $last_len); + # A custom archive's WriteInt() length has four little-endian bytes. + my $length_size = length(pack('V', 0)); + # The length is preceded by WriteInt()'s one-byte sign field. + my $block_header_size = 1 + $length_size; + + # + # Compressed data in a custom archive is stored in length-prefixed blocks. + # The block header preceding a frame is independent of the frame magic. + # Unlike a directory archive, simply truncating a custom archive can make + # archive parsing fail before the decompressor sees the truncated frame. + # Remove a byte from the final block and update its length, so pg_restore + # reaches the compression library instead of rejecting the archive framing. + # When $empty is true, remove all data blocks while retaining the + # zero-length terminator. + # + + open my $fh, '+<', $path or die "could not open $path: $!"; + binmode $fh; + local $/; + $data = <$fh>; + + $pos = index($data, $magic); + die "compressed frame magic not found in $path" + if $pos < $block_header_size; + $pos -= $block_header_size; + $first_lenpos = $pos; + + while (1) + { + my $len; + + die "unexpected custom block length encoding in $path" + if substr($data, $pos, 1) ne "\0"; + $len = unpack('V', substr($data, $pos + 1, $length_size)); + last if $len == 0; + + die "invalid custom block length in $path" + if $pos + $block_header_size + $len > length($data); + $last_lenpos = $pos; + $last_len = $len; + $pos += $block_header_size + $len; + } + + if ($empty) + { + substr($data, $first_lenpos + 1, $length_size, pack('V', 0)); + substr($data, $first_lenpos + $block_header_size, + $pos - $first_lenpos, ''); + } + else + { + die "invalid final custom block length in $path" if $last_len < 2; + substr($data, $last_lenpos + 1, $length_size, + pack('V', $last_len - 1)); + substr($data, $last_lenpos + $block_header_size + $last_len - 1, + 1, ''); + } + + seek($fh, 0, 0) or die "could not seek $path: $!"; + truncate($fh, 0) or die "could not truncate $path: $!"; + print {$fh} $data or die "could not write $path: $!"; + close $fh or die "could not close $path: $!"; +} + my %pgdump_runs = ( compression_none_custom => { test_key => 'compression', @@ -631,6 +701,61 @@ foreach my $run (sort keys %pgdump_runs) } } +for my $method (qw(lz4 zstd)) +{ + my $supported = $method eq 'lz4' ? $supports_lz4 : $supports_zstd; + my $extension = $method eq 'lz4' ? 'lz4' : 'zst'; + my $magic = $method eq 'lz4' ? pack('H*', '04224d18') : + pack('H*', '28b52ffd'); + my $custom_path = "$tempdir/compression_${method}_custom.dump"; + my $custom_bad_path = "$tempdir/${method}_custom_bad.dump"; + my ($directory_path) = glob("$tempdir/compression_${method}_dir/*.dat.$extension"); + + SKIP: + { + skip "$method compression not supported by this build", 2 + if !$supported; + + copy($custom_path, $custom_bad_path) + or die "could not copy $custom_path: $!"; + truncate_custom_compressed_data($custom_bad_path, $magic, 0); + $node->command_fails( + [ 'pg_restore', '--file' => "$tempdir/${method}_custom_bad.sql", + $custom_bad_path ], + "$method custom archive with truncated compressed data fails"); + + die "could not find compressed data file for $method" + if !defined($directory_path); + my $directory_size = -s $directory_path; + truncate($directory_path, $directory_size - 1) + or die "could not truncate $directory_path: $!"; + $node->command_fails( + [ 'pg_restore', '--file' => "$tempdir/${method}_directory_bad.sql", + "$tempdir/compression_${method}_dir" ], + "$method directory archive with truncated compressed data fails"); + } +} + +SKIP: +{ + # Zstd needs empty-input processing at end-of-input to flush buffered + # output. LZ4 has no equivalent processing; its completion state remains + # false for an empty stream. + skip "zstd compression not supported by this build", 1 if !$supports_zstd; + + my $custom_path = "$tempdir/compression_zstd_custom.dump"; + my $custom_bad_path = "$tempdir/zstd_custom_empty.dump"; + + copy($custom_path, $custom_bad_path) + or die "could not copy $custom_path: $!"; + truncate_custom_compressed_data($custom_bad_path, + pack('H*', '28b52ffd'), 1); + $node->command_fails( + [ 'pg_restore', '--file' => "$tempdir/zstd_custom_empty.sql", + $custom_bad_path ], + "zstd custom archive with empty compressed data fails"); +} + ######################################### # Stop the database instance, which will be removed at the end of the tests. -- 2.50.1 (Apple Git-155)