#!/usr/bin/env python3
"""Minimal reproducer for MDEV-41742 (_ma_rec_pack wild-pointer crash in
Aria dynamic-record write, triggered via UNION DISTINCT over a two-table
join with a binary-protocol prepared statement string parameter).

Reduced from the original: schema trimmed to only the columns referenced
by the query (t1.c1 for the join; t2.c1/c2/c7 for the join + select list),
both tables' PRIMARY KEY/UNIQUE KEY dropped, and the parameter shrunk from
32869 to 16384 bytes (the ticket's own stated minimum for utf8mb4:
16384 * mbmaxlen(4) = 65536 > 65535).

Untouched per the ticket's own ablation (verified necessary):
  - UNION DISTINCT (not ALL)
  - exactly 6 projection columns (5 does not crash)
  - two-table inner join (drives the join-cache code path)
  - 1 row per table

Usage: python3 repro.py -H <host> -P <port> -u root [-d dbname]
"""
import argparse
import socket
import struct
import sys
import time

import mysql.connector

SETUP_SQL = [
    "SET SESSION sql_mode=''",
    # Forces the UNION's dedup temp table off HEAP (MDEV-38975 added HEAP
    # blob support, so without this it never converts to Aria).
    "SET SESSION tmp_memory_table_size=1024",
    "SET SESSION max_heap_table_size=1024",
    """CREATE TABLE t1 (
  c1 int NOT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",
    """CREATE TABLE t2 (
  c1 int DEFAULT NULL,
  c2 double DEFAULT NULL,
  c7 text DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4""",
    "INSERT INTO t1 (c1) VALUES (263091433)",
    "INSERT INTO t2 (c1, c2, c7) VALUES (0, 97.6, 'we4g')",
]

PREPARE_SQL = (
    "select ref_1.c7, ?, 1, ref_1.c2, ref_0.c1, ref_1.c7 "
    "from t1 ref_0 inner join t2 ref_1 on (ref_0.c1 >= ref_1.c1) "
    "union (select 1,2,3,4,5,6)"
)

# Ticket's own stated minimum for utf8mb4: 16384 * mbmaxlen(4) = 65536 > 65535.
PARAM_VALUE = b"E" * 16384


def lenenc(data: bytes) -> bytes:
    n = len(data)
    if n < 0xFB:
        return bytes([n]) + data
    if n <= 0xFFFF:
        return b"\xfc" + struct.pack("<H", n) + data
    if n <= 0xFFFFFF:
        return b"\xfd" + struct.pack("<I", n)[:3] + data
    return b"\xfe" + struct.pack("<Q", n) + data


def send_packet(sock, payload, seq=0):
    sock.sendall(struct.pack("<I", len(payload))[:3] + bytes([seq & 0xFF]) + payload)


def recv_exact(sock, n):
    data = bytearray()
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise ConnectionError("server closed the connection")
        data.extend(chunk)
    return bytes(data)


def read_packet(sock):
    hdr = recv_exact(sock, 4)
    length = hdr[0] | (hdr[1] << 8) | (hdr[2] << 16)
    return recv_exact(sock, length) if length else b""


def server_alive(args):
    try:
        c = mysql.connector.connect(host=args.host, port=args.port,
                                    user=args.user, password=args.password,
                                    connection_timeout=2, use_pure=True)
        c.close()
        return True
    except Exception:
        return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-H", "--host", default="127.0.0.1")
    ap.add_argument("-P", "--port", type=int, default=13309)
    ap.add_argument("-u", "--user", default="root")
    ap.add_argument("-p", "--password", default="")
    ap.add_argument("-d", "--database", default="mdev41742_repro")
    ap.add_argument("--timeout", type=int, default=20)
    args = ap.parse_args()

    conn = mysql.connector.connect(host=args.host, port=args.port,
                                   user=args.user, password=args.password,
                                   use_pure=True, autocommit=True)
    cur = conn.cursor()
    cur.execute(f"DROP DATABASE IF EXISTS `{args.database}`")
    cur.execute(f"CREATE DATABASE `{args.database}`")
    cur.execute(f"USE `{args.database}`")
    for sql in SETUP_SQL:
        cur.execute(sql)
    cur.close()
    print("[*] setup applied: 2 tables, 1 row each")

    conn.database = args.database
    sock = conn._socket.sock
    sock.settimeout(args.timeout)

    send_packet(sock, b"\x16" + PREPARE_SQL.encode("utf-8"))
    first = read_packet(sock)
    if first[0] == 0xFF:
        errno = struct.unpack_from("<H", first, 1)[0]
        print(f"[-] PREPARE failed: errno={errno}")
        sys.exit(1)
    stmt_id = struct.unpack_from("<I", first, 1)[0]
    num_params = struct.unpack_from("<H", first, 7)[0]
    num_cols = struct.unpack_from("<H", first, 5)[0]
    for _ in range(num_params):
        read_packet(sock)
    if num_params:
        read_packet(sock)
    for _ in range(num_cols):
        read_packet(sock)
    if num_cols:
        read_packet(sock)
    print(f"[*] prepare ok: stmt_id={stmt_id} num_params={num_params} "
          f"num_columns={num_cols}")
    if num_params != 1:
        print("[-] expected 1 parameter")
        sys.exit(1)

    # flags=0x00 -- ticket states COM_STMT_EXECUTE flags are irrelevant.
    payload = (b"\x17" + struct.pack("<I", stmt_id) + b"\x00" +
               struct.pack("<I", 1) + b"\x00" + b"\x01" +
               b"\xfe\x00" + lenenc(PARAM_VALUE))
    print(f"[*] sending EXECUTE: param0=STRING(0xfe) {len(PARAM_VALUE)} bytes")
    send_packet(sock, payload)

    try:
        resp = read_packet(sock)
        if resp and resp[0] == 0xFF:
            errno = struct.unpack_from("<H", resp, 1)[0]
            print(f"[-] server replied ERR errno={errno} (bug NOT triggered)")
        else:
            print("[-] server survived EXECUTE (bug NOT triggered)")
        return
    except (socket.timeout, ConnectionError, OSError) as exc:
        print(f"[!] connection lost during EXECUTE: {exc}")

    time.sleep(1.5)
    if server_alive(args):
        print("[-] connection lost but server alive (bug NOT confirmed)")
    else:
        print("[!!!] BUG CONFIRMED: server appears to be down "
              "(expected ASan failure at _ma_rec_pack@ma_dynrec.c:1023)")


if __name__ == "__main__":
    main()