packmate/stress/network/verify_ledger.py

import argparse
import hashlib
import json
import struct
import sys
from pathlib import Path


PAYLOAD_MARKER = b"PMSTRESS"
PAD_CYCLE = bytes(range(251))


def build_payload(run_id: int, connection: int, shard: int, sequence: int, payload_bytes: int) -> bytes:
    prefix = struct.pack("!QQII", run_id, connection, shard, sequence)
    padding_size = payload_bytes - len(prefix)
    start = (connection + sequence) % len(PAD_CYCLE)
    rotated = PAD_CYCLE[start:] + PAD_CYCLE[:start]
    padding = bytearray((rotated * ((padding_size + len(PAD_CYCLE) - 1) // len(PAD_CYCLE)))[:padding_size])
    if sequence == 0:
        padding[:len(PAYLOAD_MARKER)] = PAYLOAD_MARKER
    return prefix + padding


def expected_digest(
    run_id: int,
    connection: int,
    shard: int,
    frames: int,
    payload_bytes: int,
) -> str:
    digest = hashlib.md5()
    for sequence in range(frames):
        payload = build_payload(run_id, connection, shard, sequence, payload_bytes)
        digest.update(struct.pack("!I", len(payload)))
        digest.update(payload)
    return digest.hexdigest()


def expected_shard(connection: int, streams: int, shards: int) -> int:
    base, remainder = divmod(streams, shards)
    larger_range = (base + 1) * remainder
    if connection < larger_range:
        return connection // (base + 1)
    if base == 0:
        raise ValueError(f"connection {connection} is outside non-empty shards")
    return remainder + (connection - larger_range) // base


def verify(args: argparse.Namespace) -> dict:
    errors: list[str] = []
    seen_connections: set[int] = set()
    rows = 0
    valid_rows = 0
    expected_length = args.frames * (args.payload_bytes + 4)

    def fail(message: str) -> None:
        if len(errors) < 100:
            errors.append(message)

    for raw_line in sys.stdin:
        line = raw_line.rstrip("\r\n")
        if not line:
            continue
        rows += 1
        stream_id = "?"
        try:
            (
                stream_id_text,
                marker_hex,
                incoming_length_text,
                incoming_hash,
                outgoing_length_text,
                outgoing_hash,
                ips_valid_text,
                metadata_valid_text,
            ) = line.split("\t")
            stream_id = int(stream_id_text)
            marker = bytes.fromhex(marker_hex)
            if len(marker) != 24:
                raise ValueError(f"marker has {len(marker)} bytes")
            run_id, connection, shard, sequence = struct.unpack("!QQII", marker)
            if run_id != args.run_id or sequence != 0:
                raise ValueError(f"marker run={run_id} sequence={sequence}")
            if connection >= args.total_streams:
                raise ValueError(f"connection {connection} is out of range")
            required_shard = expected_shard(connection, args.total_streams, args.shards)
            if shard != required_shard:
                raise ValueError(f"connection {connection} has shard={shard}, expected={required_shard}")
            if connection in seen_connections:
                raise ValueError(f"connection {connection} appears in more than one captured stream")
            incoming_length = int(incoming_length_text)
            outgoing_length = int(outgoing_length_text)
            if incoming_length != expected_length or outgoing_length != expected_length:
                raise ValueError(
                    f"payload lengths incoming={incoming_length} outgoing={outgoing_length} "
                    f"expected={expected_length}"
                )
            digest = expected_digest(run_id, connection, shard, args.frames, args.payload_bytes)
            if incoming_hash != digest or outgoing_hash != digest:
                raise ValueError(
                    f"payload hash incoming={incoming_hash} outgoing={outgoing_hash} expected={digest}"
                )
            if ips_valid_text != "t":
                raise ValueError("packet IP directions are invalid")
            if metadata_valid_text != "t":
                raise ValueError("stream packet count, byte size, service, protocol or timestamps are invalid")
            seen_connections.add(connection)
            valid_rows += 1
        except (ValueError, struct.error) as exc:
            fail(f"row={rows} stream={stream_id}: {exc}")

    if rows != args.expected_streams:
        fail(f"captured ledger contains {rows} streams, expected {args.expected_streams}")
    if not args.allow_subset and len(seen_connections) != args.total_streams:
        missing = []
        for connection in range(args.total_streams):
            if connection not in seen_connections:
                missing.append(connection)
                if len(missing) == 20:
                    break
        fail(
            f"captured unique connections={len(seen_connections)}, expected={args.total_streams}, "
            f"sampleMissing={missing}"
        )

    result = {
        "runId": args.run_id,
        "ledgerRows": rows,
        "validRows": valid_rows,
        "uniqueConnections": len(seen_connections),
        "expectedStreams": args.expected_streams,
        "totalOfferedStreams": args.total_streams,
        "allowSubset": args.allow_subset,
        "failedRows": rows - valid_rows,
        "sampleErrors": errors,
    }
    Path(args.report).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result), flush=True)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Verify every captured connection from a PostgreSQL TSV ledger")
    parser.add_argument("--run-id", type=int, required=True)
    parser.add_argument("--total-streams", type=int, required=True)
    parser.add_argument("--expected-streams", type=int, required=True)
    parser.add_argument("--shards", type=int, required=True)
    parser.add_argument("--frames", type=int, required=True)
    parser.add_argument("--payload-bytes", type=int, required=True)
    parser.add_argument("--allow-subset", action="store_true")
    parser.add_argument("--report", required=True)
    args = parser.parse_args()
    if min(
        args.run_id,
        args.total_streams,
        args.expected_streams,
        args.shards,
        args.frames,
        args.payload_bytes,
    ) <= 0:
        parser.error("all numeric arguments must be positive")
    if args.expected_streams > args.total_streams:
        parser.error("expected streams cannot exceed total offered streams")
    return args


if __name__ == "__main__":
    arguments = parse_args()
    result = verify(arguments)
    raise SystemExit(0 if not result["sampleErrors"] else 1)