workdir=$(mktemp -d /tmp/pgdump-trunc.XXXXXX)
bindir=/usr/local/pgsql/bin
port=55489

truncate_custom_block() {
  perl -e '
    use strict;
    use warnings;

    my ($path, $method) = @ARGV;

    open my $fh, "+<", $path or die "could not open $path: $!";
    binmode $fh;
    local $/;
    my $data = <$fh>;

    my $pos = -1;

    if ($method eq "zstd")
    {
        $pos = index($data, pack("H*", "28b52ffd"));
    }
    elsif ($method eq "lz4")
    {
        $pos = index($data, pack("H*", "04224d18"));
    }
    elsif ($method eq "gzip")
    {
        for (my $i = 0; $i + 1 < length($data); $i++)
        {
            next if substr($data, $i, 1) ne "\x78";

            my $flg = ord(substr($data, $i + 1, 1));
            if (((0x78 << 8) + $flg) % 31 == 0)
            {
                $pos = $i;
                last;
            }
        }
    }
    else
    {
        die "unknown compression method: $method\n";
    }

    die "compressed-frame header not found\n" if $pos < 5;

    my $lenpos = $pos - 5;
    die "unexpected custom block length encoding\n"
      if substr($data, $lenpos, 1) ne "\0";

    my $len = unpack("V", substr($data, $lenpos + 1, 4));
    die "invalid custom block length\n"
      if $len < 2 || $pos + $len > length($data);

    substr($data, $lenpos + 1, 4, pack("V", $len - 1));
    substr($data, $pos + $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: $!";
  ' "$1" "$2"
}

"$bindir/initdb" -D "$workdir/data" -N --no-sync
"$bindir/pg_ctl" -D "$workdir/data" -l "$workdir/server.log" \
  -o "-p $port -k $workdir -c listen_addresses=" start

"$bindir/psql" -h "$workdir" -p "$port" -d postgres -c \
  "CREATE TABLE t AS
   SELECT i, repeat(md5(i::text), 5) AS data
   FROM generate_series(1, 400) AS g(i);"

for method in zstd lz4 gzip; do
  "$bindir/pg_dump" -h "$workdir" -p "$port" \
    -Fc --compress="$method" -f "$workdir/$method.dump" postgres

  cp "$workdir/$method.dump" "$workdir/$method-custom-bad.dump"
  truncate_custom_block "$workdir/$method-custom-bad.dump" "$method"

  "$bindir/pg_restore" -f "$workdir/$method-custom.sql" \
    "$workdir/$method-custom-bad.dump"
  printf 'custom %s: exit status %s\n' "$method" "$?"

  "$bindir/pg_dump" -h "$workdir" -p "$port" \
    -Fd -j 1 --compress="$method" -f "$workdir/$method-dir" postgres

  cp -R "$workdir/$method-dir" "$workdir/$method-dir-bad"
  archive=$(find "$workdir/$method-dir-bad" -name '*.dat.*' | head -n 1)
  truncate -s -1 "$archive"

  "$bindir/pg_restore" -f "$workdir/$method-dir.sql" \
    "$workdir/$method-dir-bad"
  printf 'directory %s: exit status %s\n' "$method" "$?"
done

"$bindir/pg_ctl" -D "$workdir/data" stop -m fast