niko_trust/tools/reference/pow_reference.py

#!/usr/bin/env python3
"""Reference implementation of the relay's proof-of-work admission control.

Companion to docs/POW.md and the Go package internal/pow, with the same two
purposes as tce_reference.py: generate the frozen vectors in
testdata/vectors/pow_vectors.json, and be an independent implementation so
that agreement means "two readings of the specification coincide".

Scheme:

    hash = BLAKE3_keyed(key, DOMAIN || target || counter_be)

valid when `hash` has at least `difficulty` leading zero bits.

Every message here is at most len(DOMAIN)+32+4 = 57 bytes, so the reference
implements exactly the BLAKE3 rule set for single-block inputs: key words as
the chaining value, one compression call over one padded 64-byte block, with
CHUNK_START | CHUNK_END | KEYED_HASH flags. That is complete for this input
class, not an approximation.

Run:  python3 tools/reference/pow_reference.py --emit testdata/vectors
"""

from __future__ import annotations

import argparse
import json
import struct
from pathlib import Path

DOMAIN = b"trust.n1ko.dev/pow/1"
KEY_SIZE = 32
TARGET_SIZE = 32
MAX_DIFFICULTY = 30

# BLAKE3 flags for a keyed hash whose entire input is one block: the chunk
# starts and ends here, the whole input is the root node, and the hash is
# keyed.
FLAG_CHUNK_START = 1 << 0
FLAG_CHUNK_END = 1 << 1
FLAG_ROOT = 1 << 3
FLAG_KEYED_HASH = 1 << 4
SINGLE_BLOCK_FLAGS = FLAG_CHUNK_START | FLAG_CHUNK_END | FLAG_ROOT | FLAG_KEYED_HASH

IV = (
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
    0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
)
MSG_PERMUTATION = (2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8)


def _rotr(x: int, n: int) -> int:
    return ((x >> n) | (x << (32 - n))) & 0xFFFFFFFF


def _g(s: list[int], a: int, b: int, c: int, d: int, mx: int, my: int) -> None:
    s[a] = (s[a] + s[b] + mx) & 0xFFFFFFFF
    s[d] = _rotr(s[d] ^ s[a], 16)
    s[c] = (s[c] + s[d]) & 0xFFFFFFFF
    s[b] = _rotr(s[b] ^ s[c], 12)
    s[a] = (s[a] + s[b] + my) & 0xFFFFFFFF
    s[d] = _rotr(s[d] ^ s[a], 8)
    s[c] = (s[c] + s[d]) & 0xFFFFFFFF
    s[b] = _rotr(s[b] ^ s[c], 7)


def _compress(cv: tuple[int, ...], block_words: list[int], block_len: int, flags: int) -> list[int]:
    counter = 0  # first and only chunk
    state = list(cv) + list(IV[:4]) + [
        counter & 0xFFFFFFFF,
        (counter >> 32) & 0xFFFFFFFF,
        block_len,
        flags,
    ]
    m = list(block_words)
    for r in range(7):
        if r > 0:
            m = [m[MSG_PERMUTATION[i]] for i in range(16)]
        _g(state, 0, 4, 8, 12, m[0], m[1])
        _g(state, 1, 5, 9, 13, m[2], m[3])
        _g(state, 2, 6, 10, 14, m[4], m[5])
        _g(state, 3, 7, 11, 15, m[6], m[7])
        _g(state, 0, 5, 10, 15, m[8], m[9])
        _g(state, 1, 6, 11, 12, m[10], m[11])
        _g(state, 2, 7, 8, 13, m[12], m[13])
        _g(state, 3, 4, 9, 14, m[14], m[15])
    # Final transform: the first half becomes the new chaining value, the
    # second half the block output. For a single-chunk root input the digest
    # is the new chaining value: root_output_bytes streams all 16 words, and
    # a 32-byte digest reads exactly the first eight.
    for i in range(8):
        state[i] ^= state[i + 8]
        state[i + 8] ^= cv[i]
    return state


def pow_hash(key: bytes, target: bytes, counter: int) -> bytes:
    """BLAKE3_keyed(key, DOMAIN || target || counter_be), 32-byte digest."""
    if len(key) != KEY_SIZE:
        raise ValueError("key must be 32 bytes")
    if len(target) != TARGET_SIZE:
        raise ValueError("target must be 32 bytes")
    msg = DOMAIN + target + struct.pack(">I", counter)
    if len(msg) > 64:
        raise ValueError("reference covers single-block inputs only")

    padded = msg + bytes(64 - len(msg))
    block = list(struct.unpack("<16I", padded))
    cv = struct.unpack("<8I", key)
    words = _compress(cv, block, len(msg), SINGLE_BLOCK_FLAGS)
    return struct.pack("<8I", *words[:8])


def leading_zero_bits(digest: bytes) -> int:
    n = 0
    for byte in digest:
        if byte == 0:
            n += 8
            continue
        mask = 0x80
        while mask and not byte & mask:
            n += 1
            mask >>= 1
        break
    return n


def solve(key: bytes, target: bytes, difficulty: int) -> int:
    counter = 0
    while True:
        if leading_zero_bits(pow_hash(key, target, counter)) >= difficulty:
            return counter
        counter += 1


# --------------------------------------------------------------------------
# Vector generation
# --------------------------------------------------------------------------

def fixture(seed: int) -> bytes:
    return bytes([seed]) * KEY_SIZE


def emit(out_dir: Path) -> None:
    vectors = []
    rejects = []

    # (key_seed, target_seed or None for zeros, difficulty)
    cases = [
        (1, 2, 1),
        (1, 2, 4),
        (1, 2, 8),
        (2, 0, 8),      # zero target: the authentication binding
        (2, 3, 12),
        (4, 5, 16),
    ]
    for key_seed, target_seed, difficulty in cases:
        key = fixture(key_seed)
        target = bytes(TARGET_SIZE) if target_seed == 0 else fixture(target_seed)
        counter = solve(key, target, difficulty)
        name = "auth/zero-target" if target_seed == 0 else (
            f"put/key-seed{key_seed:02x}-target-seed{target_seed:02x}"
        )
        vectors.append({
            "name": f"{name}/diff{difficulty}",
            "key_hex": key.hex(),
            "target_hex": target.hex(),
            "difficulty": difficulty,
            "counter": counter,
            "sum_hex": pow_hash(key, target, counter).hex(),
            "leading_zero_bits": leading_zero_bits(pow_hash(key, target, counter)),
        })

    # Rejects derived from the first vector: off-by-one counter and wrong key.
    base = cases[2]
    key = fixture(base[0])
    target = fixture(base[1])
    good = solve(key, target, base[2])
    bad_counter = good + 1
    while leading_zero_bits(pow_hash(key, target, bad_counter)) >= base[2]:
        bad_counter += 1
    rejects.append({
        "name": "counter-misses-threshold",
        "key_hex": key.hex(),
        "target_hex": target.hex(),
        "difficulty": base[2],
        "counter": bad_counter,
        "reason": "leading_zero_bits < difficulty",
    })
    rejects.append({
        "name": "wrong-key",
        "key_hex": fixture(0xFF).hex(),
        "target_hex": target.hex(),
        "difficulty": base[2],
        "counter": good,
        "reason": "hash under a different key does not meet the target",
    })

    config_rejects = [
        {
            "name": "difficulty-above-max",
            "difficulty": MAX_DIFFICULTY + 1,
            "reason": "configuration above MaxDifficulty is clamped/rejected",
        },
        {
            "name": "negative-difficulty",
            "difficulty": -1,
            "reason": "negative configuration is treated as disabled, never as a pass-all proof tier",
        },
    ]

    doc = {
        "domain": DOMAIN.decode(),
        "hash_spec": "BLAKE3_keyed(key, domain || target || counter_be); valid iff leading_zero_bits(sum) >= difficulty",
        "max_difficulty": MAX_DIFFICULTY,
        "vectors": vectors,
        "rejects": rejects,
        "config_rejects": config_rejects,
    }
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / "pow_vectors.json"
    path.write_text(json.dumps(doc, indent=2) + "\n")
    print(f"wrote {path} ({len(vectors)} vectors, {len(rejects)} rejects)")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--emit", default="testdata/vectors", help="output directory")
    args = ap.parse_args()
    emit(Path(args.emit))