#!/bin/bash
# Why the ZSTD_DCtx_reset() at the start of each use of the reused TOAST
# decompression context is required, and the silent wrong-data case it closes.
#
#   BIN=/path/to/postgres/install/bin bash repro-dctx-reset.sh
#
# Needs corrupt_block.py in the same directory.  Builds one table of zstd
# values; one value (id 2, 1 MB, 8 blocks) is corrupted on disk so its third
# block fails to decode mid-frame, after two good blocks.  Then, in ONE psql
# session (so one backend reuses the same decompression context throughout):
#   E*  reading the corrupt value fails with "compressed zstd data is corrupt";
#   A*  every read right after is still correct;
#   B*  the SILENT case: a prefix that stops the stream early (B1), then a
#       slice (B2) and a full read (B3) of ANOTHER value;
#   S*  the error inside a subtransaction, rolled back, then a good read;
#   M*  every good value, slice and full, in one query.
# Each value is checked by md5.  With the reset in place every A/B/S/M check is
# true.  Without it (remove the ZSTD_DCtx_reset() call and rebuild), B2 comes
# back WRONG with no error while the full reads stay correct, because
# ZSTD_decompressDCtx() starts a fresh frame but the streaming slice path
# continues from the stale mid-frame state left by B1.
set -eu
B=${BIN:?set BIN=/path/to/install/bin}
A=$(cd "$(dirname "$0")" && pwd)
D=$(mktemp -d); P=$((56400 + RANDOM % 100))
export ASAN_OPTIONS=detect_leaks=0:abort_on_error=0:log_path=$D/asan
"$B/initdb" -D $D/data -A trust --no-sync --no-data-checksums -U postgres -E SQL_ASCII --locale=C >/dev/null
cat >> $D/data/postgresql.conf <<EOF
port = $P
unix_socket_directories = '/tmp'
autovacuum = off
jit = off
EOF
start() { "$B/pg_ctl" -D $D/data -l $D/log -w -t 60 start >/dev/null; }
cleanup() { "$B/pg_ctl" -D $D/data -w -m fast stop >/dev/null 2>&1 || true; rm -rf $D; }
trap cleanup EXIT
start
q() { "$B/psql" -X -qAt -h /tmp -p $P -U postgres -d postgres "$@"; }

q -v ON_ERROR_STOP=1 <<'EOF'
CREATE TABLE t (id int PRIMARY KEY, v text COMPRESSION zstd, h text, p int, hs text);
WITH v(id, x, p) AS (
  SELECT 1, (SELECT string_agg(md5((100000 + i)::text), '') FROM generate_series(1, 31250) i), 500000
  UNION ALL
  SELECT 2, (SELECT string_agg(md5((200000 + i)::text), '') FROM generate_series(1, 31250) i), 500000
  UNION ALL
  SELECT 2 + g, (SELECT string_agg(md5((g * 1000 + i)::text) || ' row', '') FROM generate_series(1, 800) i), 16000
    FROM generate_series(1, 200) g
  UNION ALL
  SELECT 1000 + g, repeat('abc' || g, 1000), 2000 FROM generate_series(1, 20) g)
INSERT INTO t SELECT id, x, md5(x), p, md5(substr(x, p, 100)) FROM v;
EOF
echo "values: $(q -c "SELECT count(*) || ' rows, ' || count(*) FILTER (WHERE pg_column_compression(v) = 'zstd') || ' zstd, ' || count(*) FILTER (WHERE pg_column_toast_chunk_id(v) IS NOT NULL) || ' external' FROM t")"

# Dump the chunks of value 2, stop, corrupt its third block, start.
toast=$(q -c "SELECT reltoastrelid::regclass FROM pg_class WHERE relname = 't'")
q -c "COPY (SELECT chunk_seq || '|' || encode(chunk_data, 'hex') FROM $toast WHERE chunk_id::text = (SELECT pg_column_toast_chunk_id(v)::text FROM t WHERE id = 2) ORDER BY chunk_seq) TO STDOUT" > $D/chunks.txt
relfile=$D/data/$(q -c "SELECT pg_relation_filepath('$toast')")
q -c CHECKPOINT
"$B/pg_ctl" -D $D/data -w stop >/dev/null
python3 -I $A/corrupt_block.py $D/chunks.txt $relfile 2
start

# All in one session.  ON_ERROR_STOP off: the E lines are meant to fail.
q -v ON_ERROR_STOP=0 > $D/out.txt 2>&1 <<'EOF'
SELECT 'pid ' || pg_backend_pid();
SELECT 'A0 slice good ' || (md5(substr(v, p, 100)) = hs) FROM t WHERE id = 1;
SELECT 'E1 slice corrupt ' || length(substr(v, p, 100)) FROM t WHERE id = 2;
SELECT 'A1 slice good after slice error ' || (md5(substr(v, p, 100)) = hs) FROM t WHERE id = 1;
SELECT 'A2 full good after slice error ' || (md5(v) = h) FROM t WHERE id = 1;
SELECT 'E2 full corrupt ' || length(md5(v)) FROM t WHERE id = 2;
SELECT 'A3 full good after full error ' || (md5(v) = h) FROM t WHERE id = 1;
SELECT 'A4 slice good after full error ' || (md5(substr(v, p, 100)) = hs) FROM t WHERE id = 1;
SELECT 'P1 prefix of the corrupt value, before the bad block ' || (substr(v, 1, 100) = substr(md5('200001') || md5('200002') || md5('200003') || md5('200004'), 1, 100)) FROM t WHERE id = 2;
SELECT 'B1 prefix (stops early) ' || length(substr(v, 1, 100)) FROM t WHERE id = 1;
SELECT 'B2 slice of another value after an early stop ' || (md5(substr(v, p, 100)) = hs) FROM t WHERE id = 3;
SELECT 'B3 full of another value after an early stop ' || (md5(v) = h) FROM t WHERE id = 3;
BEGIN;
SAVEPOINT s;
SELECT 'S1 slice corrupt in a subtransaction ' || length(substr(v, p, 100)) FROM t WHERE id = 2;
ROLLBACK TO s;
SELECT 'S2 slice good after rollback to savepoint ' || (md5(substr(v, p, 100)) = hs) FROM t WHERE id = 1;
COMMIT;
SELECT 'M1 all good values, slice then full, ' || count(*) || ' rows: ' || bool_and(md5(substr(v, p, 100)) = hs AND md5(v) = h) FROM t WHERE id <> 2;
SELECT 'M2 prefix of every value then slice, ' || count(*) || ' rows: ' || bool_and(length(substr(v, 1, 10)) = 10 AND md5(substr(v, p, 100)) = hs) FROM t WHERE id <> 2;
EOF
cat $D/out.txt
exp_err=$(grep -c "compressed zstd data is corrupt" $D/out.txt || true)
falses=$(grep -cE "^[ABSMP][0-9].* false$" $D/out.txt || true)
others=$(grep -E "ERROR" $D/out.txt | grep -vc "compressed zstd data is corrupt" || true)
asan=$(cat $D/asan.* 2>/dev/null | grep -c "ERROR: AddressSanitizer" || true)
"$B/pg_ctl" -D $D/data -w stop >/dev/null
echo "corrupt-errors=$exp_err (expected 3) false=$falses other-errors=$others asan-reports=$asan"
if [ "$exp_err" = 3 ] && [ "$falses" = 0 ] && [ "$others" = 0 ] && [ "$asan" = 0 ] && [ $(grep -cE "^[ABSMP][0-9].* true$|^B1 .* 100$" $D/out.txt) = 12 ]; then
  echo "RESULT pass (reset present)"
else
  echo "RESULT FAIL (expected without the reset: B2 false, others still true)"
fi