From 80b0a2541b0603746978f321c2804ae4b7d8738a Mon Sep 17 00:00:00 2001 From: "Chao Li (Evan)" Date: Fri, 14 Aug 2026 11:29:21 +0800 Subject: [PATCH v8 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: Discussion: https://postgr.es/m/5962B878-C43D-4EBC-9E95-1F945CE5E586@gmail.com --- src/bin/pg_dump/compress_lz4.c | 25 ++++-- src/bin/pg_dump/compress_zstd.c | 67 +++++++++++++- src/bin/pg_dump/t/006_pg_dump_compress.pl | 103 ++++++++++++++++++++++ 3 files changed, 187 insertions(+), 8 deletions(-) diff --git a/src/bin/pg_dump/compress_lz4.c b/src/bin/pg_dump/compress_lz4.c index 500d5e16a6d..489f3abcbb9 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. @@ -159,6 +160,7 @@ ReadDataFromArchiveLZ4(ArchiveHandle *AH, CompressorState *cs) LZ4F_decompressionContext_t ctx = NULL; LZ4F_decompressOptions_t dec_opt; LZ4F_errorCode_t status; + bool frame_finished = false; memset(&dec_opt, 0, sizeof(dec_opt)); status = LZ4F_createDecompressionContext(&ctx, LZ4F_VERSION); @@ -187,12 +189,16 @@ ReadDataFromArchiveLZ4(ArchiveHandle *AH, CompressorState *cs) if (LZ4F_isError(status)) pg_fatal("could not decompress: %s", LZ4F_getErrorName(status)); + frame_finished = (status == 0); ahwrite(outbuf, 1, out_size, AH); readp += read_size; } } + if (!frame_finished) + pg_fatal("could not decompress data: compressed stream is incomplete"); + pg_free(outbuf); pg_free(readbuf); @@ -466,12 +472,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 +503,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; @@ -681,6 +691,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..ef225235ed9 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,56 @@ 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); + + 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 < 5; + $pos -= 5; + $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, 4)); + last if $len == 0; + + die "invalid custom block length in $path" + if $pos + 5 + $len > length($data); + $last_lenpos = $pos; + $last_len = $len; + $pos += 5 + $len; + } + + if ($empty) + { + substr($data, $first_lenpos + 1, 4, pack('V', 0)); + substr($data, $first_lenpos + 5, $pos - $first_lenpos, ''); + } + else + { + die "invalid final custom block length in $path" if $last_len < 2; + substr($data, $last_lenpos + 1, 4, pack('V', $last_len - 1)); + substr($data, $last_lenpos + 5 + $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 +682,58 @@ 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: +{ + 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)