From 7afbca00d52fd8512e53eede651bf9ebecb25638 Mon Sep 17 00:00:00 2001
From: "Chao Li (Evan)" <lic@highgo.com>
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 <lic@highgo.com>
Co-authored-by: Daniel Gustafsson <daniel@yesql.se>
Reviewed-by: Daniel Gustafsson <daniel@yesql.se>
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)

