#!/bin/bash
# forge-pointers.sh PREFIX OUT.tsv
# Forged on-disk TOAST pointers of the forms the series decodes.  PREFIX is
# an install prefix (holding bin/), built with zstd; use an assert-enabled
# build.  Only the method bits of va_extinfo, the method byte of the long
# form and the saved size are rewritten, with the server stopped.
# Cases, one row and one external value each:
#   zstd (long pointer): untouched; method byte 0..255 except 2; method
#     bits 0, 1, 2; saved size -1 and +1; saved size = rawsize - 4 and
#     = rawsize (the pointer then says "not compressed")
#   pglz, lz4 (plain pointer): untouched; method bits 2 and 3
# For each: verify_heapam(check_toast => true), then pg_column_compression(),
# length(), octet_length(), md5(v), substr(v, 1, 5), substr(v, 3000, 10) and
# left(v, 100), each in its own psql; md5s taken before forging tell
# "intact" from "WRONG".  A backend crash shows as CRASH.
set -eu
I=$1; OUT=$2
D=$(mktemp -d ${TMPDIR:-/tmp}/forge.XXXXXX); P=$((61000 + RANDOM % 900))
trap '"$I/bin/pg_ctl" -D $D/data -m immediate -w stop >/dev/null 2>&1 || true; rm -rf $D' EXIT
export ASAN_OPTIONS=detect_leaks=0:abort_on_error=0:log_path=$D/asan
Q() { "$I/bin/psql" -X -qAt -h $D -p $P -U postgres -d postgres -v ON_ERROR_STOP=1 "$@"; }
"$I/bin/initdb" -D $D/data -A trust --no-sync --no-data-checksums -U postgres -E UTF8 --locale=C >/dev/null
printf "port = $P\nunix_socket_directories = '$D'\nlisten_addresses = ''\nautovacuum = off\nrestart_after_crash = on\n" >> $D/data/postgresql.conf
"$I/bin/pg_ctl" -D $D/data -l $D/log -w start >/dev/null
python3 -I - > $D/cases.txt <<'PY'
c = [('z', 'ctl', 0)] + [('z', 'byte', b) for b in range(256) if b != 2]
c += [('z', 'bits', b) for b in (0, 1, 2)]
c += [('z', k, 0) for k in ('ext_minus1', 'ext_plus1', 'ext_eq_raw', 'ext_gt_raw')]
for m in 'pl':
    c += [(m, 'ctl', 0), (m, 'bits', 2), (m, 'bits', 3)]
for i, (m, k, a) in enumerate(c, 1):
    print(i, m, k, a)
PY
Q <<'SQL'
CREATE EXTENSION amcheck;
CREATE EXTENSION pageinspect;
CREATE TABLE cases (id int, m text, kind text, arg int);
CREATE TABLE h (id int, hv text, h1 text, h2 text, h3 text);
CREATE FUNCTION val(i int) RETURNS text LANGUAGE sql IMMUTABLE AS $$
  SELECT (SELECT string_agg(md5((i * 1000 + g)::text), '') FROM generate_series(1, 60) g)
         || repeat('x', 4000) $$;
CREATE TABLE fz (id int, v text COMPRESSION zstd) WITH (toast_tuple_target = 128);
CREATE TABLE fp (id int, v text COMPRESSION pglz) WITH (toast_tuple_target = 128);
CREATE TABLE fl (id int, v text COMPRESSION lz4) WITH (toast_tuple_target = 128);
SQL
Q -c "\\copy cases FROM '$D/cases.txt' WITH (DELIMITER ' ')"
for m in z p l; do
  Q -c "INSERT INTO f$m SELECT id, val(id) FROM cases WHERE m = '$m' ORDER BY id"
  Q -c "INSERT INTO h SELECT id, md5(v), md5(substr(v, 1, 5)), md5(substr(v, 3000, 10)), md5(left(v, 100)) FROM f$m"
  # tuple data = int4 id, then the pointer: 0x01, vartag, va_rawsize,
  # va_extinfo, va_valueid, va_toastrelid, [method byte]
  Q -c "COPY (SELECT '$m', pg_relation_filepath('f$m'), b * 8192 + lp_off + t_hoff,
               encode(substr(t_data, 1, 4), 'hex'), get_byte(t_data, 4), get_byte(t_data, 5),
               encode(substr(t_data, 7, 8), 'hex')
        FROM generate_series(0, pg_relation_size('f$m') / 8192 - 1) b,
             heap_page_items(get_raw_page('f$m', b::int)) WHERE lp_len > 0) TO STDOUT" >> $D/where.txt
done
Q -c CHECKPOINT
"$I/bin/pg_ctl" -D $D/data -m fast -w stop >/dev/null
python3 -I - $D/cases.txt $D/where.txt $D/data <<'PY'
import struct, sys
cases = {}
for line in open(sys.argv[1]):
    i, m, k, a = line.split(); cases[int(i)] = (m, k, int(a))
for line in open(sys.argv[2]):
    m, path, off, idhex, b4, tag, rawext = line.split()
    i = struct.unpack('<i', bytes.fromhex(idhex))[0]
    cm, k, a = cases[i]
    assert cm == m and b4 == '1' and int(tag) == (19 if m == 'z' else 18), line
    base = int(off) + 4 + 2
    raw, ext = struct.unpack('<iI', bytes.fromhex(rawext))
    size, bits = ext & 0x3FFFFFFF, ext >> 30
    if k == 'ctl':
        continue
    if k == 'byte':
        at, new = base + 16, bytes([a])
    elif k == 'bits':
        at, new = base + 4, struct.pack('<I', size | (a << 30))
    else:
        n = {'ext_minus1': size - 1, 'ext_plus1': size + 1,
             'ext_eq_raw': raw - 4, 'ext_gt_raw': raw}[k]
        at, new = base + 4, struct.pack('<I', n | (bits << 30))
    with open(sys.argv[3] + '/' + path, 'r+b') as f:
        f.seek(at); f.write(new)
PY
"$I/bin/pg_ctl" -D $D/data -l $D/log -w start >/dev/null
# wait until the server takes connections (a crash in a step restarts it)
up() { until "$I/bin/psql" -X -qAt -h $D -p $P -U postgres -d postgres -c 'SELECT 1' >/dev/null 2>&1; do sleep 0.2; done; }
for m in z p l; do
  up
  Q -c "COPY (SELECT x.id, regexp_replace(v.msg, 'toast value \\d+', 'toast value N')
               FROM verify_heapam('f$m', check_toast => true) v
               JOIN (SELECT id, ctid FROM f$m) x
                 ON (x.ctid::text::point)[0] = v.blkno AND (x.ctid::text::point)[1] = v.offnum
             ORDER BY 1, 2) TO STDOUT" 2>&1 || echo "0	amcheck on f$m: ERROR"
done > $D/amcheck.txt
probe() {  # one SQL in its own psql; last line, the ERROR, or CRASH
  local o
  up    # a crash in an earlier step leaves the server in recovery for a moment
  o=$("$I/bin/psql" -X -qAt -h $D -p $P -U postgres -d postgres -c "$1" 2>&1) || true
  if echo "$o" | grep -q -E 'server closed the connection|connection to server'; then
    echo CRASH
  elif echo "$o" | grep -q 'ERROR:'; then echo "$o" | grep -m1 'ERROR:' | sed -E 's/.*ERROR: +/ERROR: /; s/[0-9]{5,}/N/g'
  else echo "$o" | tail -1; fi
}
printf 'id\tmethod\tforged\targ\tamcheck\tpg_column_compression\tlength\toctet_length\tmd5\tsubstr_1_5\tsubstr_3000_10\tleft_100\n' > $OUT
while read i m k a; do
  am=$(awk -F'\t' -v i=$i '$1 == i {print $2}' $D/amcheck.txt | paste -sd';' -); row="$i	$m	$k	$a	${am:--}"
  for e in "coalesce(pg_column_compression(v), 'NULL')" "length(v)::text" "octet_length(v)::text" \
           "CASE WHEN md5(v) = hv THEN 'intact' ELSE 'WRONG' END" \
           "CASE WHEN md5(substr(v, 1, 5)) = h1 THEN 'intact' ELSE 'WRONG' END" \
           "CASE WHEN md5(substr(v, 3000, 10)) = h2 THEN 'intact' ELSE 'WRONG' END" \
           "CASE WHEN md5(left(v, 100)) = h3 THEN 'intact' ELSE 'WRONG' END"; do
    row="$row	$(probe "SELECT $e FROM f$m JOIN h USING (id) WHERE id = $i")"
  done
  echo "$row" >> $OUT
done < $D/cases.txt
echo "cases: $(($(wc -l < $OUT) - 1)); backend crashes: $(grep -c CRASH $OUT || true);" \
     "ASan reports: $(ls $D/asan* 2>/dev/null | wc -l); TRAP lines: $(grep -c TRAP $D/log || true);" \
     "verify_heapam runs that failed: $(grep -c 'amcheck on f' $D/amcheck.txt || true)"