UmbraTrans/umbratrans_web/tests/test_node_pairing.py

"""Pairing: a node exchanges a short code for a credential, and the code is spent.

This is the mechanism that replaces reading a credential out of a 0600 file and typing
it into a terminal on another machine. Those tests assert what must be true of it, and
most of them are about ways the flow could be *abused* rather than about it working:

- **Single use.** Redeeming twice yields one node and one credential. A code that
  survived would be a standing credential on a machine an operator walked away from.
- **Bound to a node name.** A code read out of a terminal buffer cannot pair a machine
  called something else. Without this, "confirm it is the machine" is decoration.
- **The code is never a credential.** Verified against the live node surface, not just
  against the pair endpoint - the property that actually matters.
- **Indistinguishable failures.** A wrong code and a code that never existed must be
  indistinguishable, or the response is an oracle and the rate limiter is decoration.
- **Rate limited before lookup.** Charged per source address ahead of any code lookup,
  so guessing across many codes costs the same as guessing at one.

Each test that asserts a refusal also asserts the row state afterwards, because "the
call failed" and "the call failed *and left nothing behind*" are different guarantees.

Skipped, with the reason, when no disposable database is configured.
"""

from __future__ import annotations

import os
import uuid
from datetime import timedelta

import pytest

from umbratrans_web.compute import NODE_SCOPES, PAIRING_ATTEMPT_BUDGET, ComputeRepository

_SCRATCH = os.environ.get("UMBRATRANS_TEST_DSN", "").strip()
_DISPOSABLE = os.environ.get("UMBRATRANS_TEST_DATABASE_IS_DISPOSABLE") == "1"

pytestmark = pytest.mark.skipif(
    not (_SCRATCH and _DISPOSABLE),
    reason="needs UMBRATRANS_TEST_DSN pointing at a disposable database",
)


@pytest.fixture
def pairing_world():
    """A repository and an organization to pair into.

    Rows are keyed to random ids and left behind rather than deleted. `node_pairing_codes`
    guards its redemption columns, so a spent code cannot be tidied away - which is the
    behaviour under test, and would also make a deleting teardown impossible.
    """

    # Constructed directly from the DSN, as test_compute.py does, rather than through
    # `load_settings`: the settings loader wants a full production environment, and a
    # test that has to fabricate one is a test that can silently stop testing.
    repository = ComputeRepository(_SCRATCH, min_size=1, max_size=4, passwords=None)
    organization = uuid.uuid4()
    with repository._connection() as connection, repository._tx(connection):
        connection.execute(
            "INSERT INTO public.organizations (id, name, slug) VALUES (%s, %s, %s)",
            (organization, f"Pairing {organization.hex[:8]}", f"pair-{organization.hex[:8]}"),
        )
    try:
        yield repository, organization
    finally:
        repository.close()


def _force_expire(repository, code_id) -> None:
    """Make a code look expired, the way the passage of time would.

    Done by moving ``created_at`` back rather than ``expires_at`` forward. The table
    carries a CHECK that ``expires_at > created_at``, and that CHECK is correct: a code
    cannot be born already expired, because then a caller would receive a code it could
    never redeem. Ageing the row is therefore only possible from behind.

    The reuse guard permits this because ``created_at`` is not identity and not
    redemption; it does still refuse a spent code changing at all, which
    `test_a_spent_code_cannot_be_revived_by_clearing_used_at` proves by writing
    directly rather than through here.
    """

    # `expires_at` is absolute, so ageing `created_at` does not expire anything - the
    # code would still be live for its full minute. The expiry has to move.
    #
    # The table carries a CHECK that `expires_at > created_at`, and the CHECK is right: a
    # code must not be born already expired, or a caller would be handed one it could
    # never redeem. So both ends move together - created_at back, expires_at with it -
    # which leaves the row internally consistent and genuinely in the past.
    with repository._connection() as connection, repository._tx(connection):
        connection.execute(
            "UPDATE public.node_pairing_codes "
            "   SET created_at = now() - interval '2 hours', "
            "       expires_at = now() - interval '1 hour' "
            " WHERE id = %s",
            (code_id,),
        )


def _mint(repository, organization, name: str, **kwargs) -> dict:
    return repository.create_pairing_code(organization, name, **kwargs)


# -- the happy path -----------------------------------------------------------


def test_a_code_is_six_digits(pairing_world) -> None:
    """Length is part of the contract: it is what the operator types.

    Asserted as a shape rather than a specific value, so a change to the alphabet or
    the length cannot pass unnoticed.
    """

    repository, organization = pairing_world
    minted = _mint(repository, organization, f"shape-{uuid.uuid4().hex[:8]}")
    assert len(minted["code"]) == 6
    assert minted["code"].isdigit()


def test_redeeming_mints_a_node_and_a_credential(pairing_world) -> None:
    """The exchange itself: a credential exists afterwards, and it works."""

    repository, organization = pairing_world
    name = f"happy-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    result = repository.redeem_pairing_code(minted["code"], name, agent_version="test/1")

    assert result["ok"] is True, result
    assert result["key"], "no key was issued"
    assert result["key"].startswith(result["key_prefix"])

    resolved = repository.resolve_node_credential(result["key"])
    assert resolved is not None, "the issued credential does not resolve"
    assert resolved[0] == uuid.UUID(str(result["node_id"]))
    # The scopes stored are the scopes the credential actually resolves to, so this
    # compares the authority rather than the response's own claim about it.
    assert sorted(resolved[2]) == sorted(result["scopes"])


def test_a_paired_node_can_do_work_and_still_cannot_read_content(pairing_world) -> None:
    """Both halves, because only one of them being right is a real failure.

    A node that cannot claim is not a compute node: the agent calls claim, complete and
    decompose, so a read-only credential leaves the entire write surface dead while every
    endpoint looks finished. A node that can read project content is a session token with
    a different name, which is the boundary `NODE_SCOPES` exists to hold.

    Written as two assertions on behaviour rather than one equality against a literal
    list, because the literal list is an implementation detail that a future scope would
    change without either of these being wrong.
    """

    repository, organization = pairing_world
    name = f"working-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    result = repository.redeem_pairing_code(minted["code"], name, agent_version="test/1")

    granted = set(result["scopes"])

    assert {"compute:claim", "compute:complete", "compute:heartbeat", "compute:read"} <= granted, (
        "a paired node cannot take or finish work, so the agent's whole write path is "
        f"unusable. It was granted {sorted(granted)}."
    )
    assert granted <= set(NODE_SCOPES), (
        f"a node credential carries scopes outside the node set: {sorted(granted - set(NODE_SCOPES))}"
    )
    assert not any(
        "project" in scope or "content" in scope or "read:all" in scope for scope in granted
    ), f"a node credential can read content: {sorted(granted)}"


def test_a_paired_node_starts_in_provisioning_not_online(pairing_world) -> None:
    """A node is online because it heartbeated, not because it was created. (I24)

    Pairing creates the machine's identity and nothing else. If this were ``online``
    before the agent had ever spoken, a node whose agent could not reach the controller
    would look available and swallow every job dispatched to it.
    """

    repository, organization = pairing_world
    name = f"provision-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    result = repository.redeem_pairing_code(minted["code"], name)

    node = repository.get_node(uuid.UUID(str(result["node_id"])))
    assert node is not None
    assert node.stored_status == "provisioning"


# -- single use ---------------------------------------------------------------


def test_a_code_cannot_be_redeemed_twice(pairing_world) -> None:
    """The property the whole design rests on.

    The second attempt must fail *and* leave the first node and credential untouched -
    a refusal that quietly issued a second credential would be worse than no refusal.
    """

    repository, organization = pairing_world
    name = f"single-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    first = repository.redeem_pairing_code(minted["code"], name)
    assert first["ok"] is True

    second = repository.redeem_pairing_code(minted["code"], name)
    assert second["ok"] is False
    assert second["code"] == "pairing_code_invalid"

    with repository._connection() as connection:
        count = connection.execute(
            "SELECT count(*) FROM public.compute_nodes WHERE name = %s", (name,)
        ).fetchone()["count"]
    assert count == 1, f"a spent code produced {count} nodes"


def test_spending_a_code_records_which_node_it_created(pairing_world) -> None:
    """Provenance, so an operator can answer 'what did this code become?'."""

    repository, organization = pairing_world
    name = f"record-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    result = repository.redeem_pairing_code(minted["code"], name)

    with repository._connection() as connection:
        row = connection.execute(
            "SELECT used_at, used_by_node_id FROM public.node_pairing_codes WHERE id = %s",
            (minted["id"],),
        ).fetchone()
    assert row["used_at"] is not None
    assert str(row["used_by_node_id"]) == str(result["node_id"])


def test_a_spent_code_cannot_be_revived_by_clearing_used_at(pairing_world) -> None:
    """The trigger, tested by attempting exactly the attack it exists to stop.

    "Modify the row to make the code usable again" is the obvious way to defeat
    single-use, and it is why the redemption columns are guarded. A guard that is only
    asserted to exist would pass here; this performs the write.
    """

    import psycopg

    repository, organization = pairing_world
    name = f"revive-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    repository.redeem_pairing_code(minted["code"], name)

    with repository._connection() as connection:
        with pytest.raises(psycopg.errors.RaiseException) as caught:
            with connection.transaction():
                connection.execute(
                    "UPDATE public.node_pairing_codes SET used_at = NULL WHERE id = %s",
                    (minted["id"],),
                )
    assert "single-use" in str(caught.value)


# -- the code is bound to a node name -----------------------------------------


def test_a_code_will_not_pair_a_differently_named_node(pairing_world) -> None:
    """The check that makes 'confirm it is the machine' mean anything.

    Without it, anyone who saw the code could pair a machine of their choosing under
    any name, and the node name an administrator typed would be decoration.
    """

    repository, organization = pairing_world
    minted = _mint(repository, organization, f"expected-{uuid.uuid4().hex[:8]}")

    result = repository.redeem_pairing_code(minted["code"], "some-other-machine")

    assert result["ok"] is False
    assert result["code"] == "pairing_name_mismatch"
    # Said plainly: an operator staring at a failure needs to be told the code is for a
    # different machine, not left guessing whether they mistyped it.
    assert "different node" in result["message"]


def test_a_name_mismatch_does_not_spend_the_code(pairing_world) -> None:
    """Otherwise a wrong name would burn a code the operator still needs.

    The obvious implementation - verify, then check the name, then spend - spends it.
    This asserts the code survives, because a mismatch is an operator error and the
    retry they will attempt with the right name has to work.
    """

    repository, organization = pairing_world
    name = f"keepalive-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    assert repository.redeem_pairing_code(minted["code"], "wrong")["ok"] is False

    retry = repository.redeem_pairing_code(minted["code"], name)
    assert retry["ok"] is True, "the code was spent by a name mismatch"


# -- the code is not a credential ---------------------------------------------


def test_the_code_does_not_authenticate_against_the_node_surface(pairing_world) -> None:
    """The property that matters most, checked on the resolver the API actually uses.

    Asserting that the pairing *code* is not accepted as a *credential* is stronger
    than checking the pair endpoint refuses it: it goes through
    `resolve_node_credential`, the same path a presented bearer key takes.
    """

    repository, organization = pairing_world
    name = f"nocred-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    repository.redeem_pairing_code(minted["code"], name)

    assert repository.resolve_node_credential(minted["code"]) is None


def test_the_plaintext_code_is_never_stored(pairing_world) -> None:
    """Hash at rest, so a database read or a support transcript yields nothing usable."""

    repository, organization = pairing_world
    name = f"hashed-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    with repository._connection() as connection:
        rows = connection.execute(
            "SELECT code_hash, code_hint FROM public.node_pairing_codes WHERE id = %s",
            (minted["id"],),
        ).fetchall()
    assert rows, "the code row is missing"
    assert minted["code"] not in rows[0]["code_hash"]
    assert minted["code"][:4] not in rows[0]["code_hash"]
    assert "argon2" in rows[0]["code_hash"]


# -- indistinguishable failures ------------------------------------------------


def test_a_wrong_code_and_a_nonexistent_one_are_indistinguishable(pairing_world) -> None:
    """Otherwise the response is an oracle.

    With the rate limiter in place, a distinguishable response is what turns a bounded
    brute force into an unbounded one: an attacker who can tell "that code exists but I
    guessed wrong" from "no such code" knows when to keep going.
    """

    repository, organization = pairing_world
    # Minted so the table is not empty: an empty table would make every comparison
    # trivially equal and the test would pass without proving anything.
    _mint(repository, organization, f"oracle-{uuid.uuid4().hex[:8]}")

    wrong = repository.redeem_pairing_code("000000", "whatever")
    absent = repository.redeem_pairing_code("999999", "whatever")

    assert wrong["code"] == absent["code"]
    assert wrong["message"] == absent["message"]


def test_an_expired_code_reads_the_same_as_a_wrong_one(pairing_world) -> None:
    """Expiry must not become a distinguishable state either."""

    repository, organization = pairing_world
    name = f"expired-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name, ttl_seconds=60)

    # Age it out rather than sleeping for a minute.
    _force_expire(repository, minted["id"])
    result = repository.redeem_pairing_code(minted["code"], name)
    assert result["ok"] is False
    assert result["code"] == "pairing_code_invalid"


def test_an_expired_code_creates_no_node(pairing_world) -> None:
    repository, organization = pairing_world
    name = f"noexpiry-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name, ttl_seconds=60)
    _force_expire(repository, minted["id"])
    repository.redeem_pairing_code(minted["code"], name)
    with repository._connection() as connection:
        count = connection.execute(
            "SELECT count(*) FROM public.compute_nodes WHERE name = %s", (name,)
        ).fetchone()["count"]
    assert count == 0


# -- rate limiting ------------------------------------------------------------


def test_repeated_guesses_are_eventually_refused(pairing_world) -> None:
    """The control that makes six digits defensible.

    Twenty guesses per quarter-hour against a million-code space needs about 21 years
    to exhaust. Without this the code length would be security theatre.
    """

    repository, organization = pairing_world
    address = f"10.0.0.{uuid.uuid4().int % 250 + 1}"
    # A real live code exists throughout, so the budget is being tested against a
    # target that would otherwise eventually be guessed rather than against nothing.
    _mint(repository, organization, f"brute-{uuid.uuid4().hex[:8]}")

    results = [
        repository.redeem_pairing_code("000000", "whatever", source_address=address)
        for _ in range(PAIRING_ATTEMPT_BUDGET + 2)
    ]
    assert results[-1]["code"] == "pairing_rate_limited", results[-1]
    assert results[0]["code"] != "pairing_rate_limited"


def test_the_limit_is_charged_before_any_code_is_looked_up(pairing_world) -> None:
    """A correct code is still refused once the address is over budget.

    This is the ordering that matters. Charging after verification would let an attacker
    use the *correct* code as an oracle for whether their guess was right, and would
    mean spraying costs nothing until one spray hits.
    """

    repository, organization = pairing_world
    address = f"10.1.0.{uuid.uuid4().int % 250 + 1}"
    name = f"budget-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    for _ in range(PAIRING_ATTEMPT_BUDGET + 1):
        repository.redeem_pairing_code("000000", "whatever", source_address=address)

    result = repository.redeem_pairing_code(minted["code"], name, source_address=address)
    assert result["ok"] is False
    assert result["code"] == "pairing_rate_limited"

    with repository._connection() as connection:
        count = connection.execute(
            "SELECT count(*) FROM public.compute_nodes WHERE name = %s", (name,)
        ).fetchone()["count"]
    assert count == 0, "an over-budget address still created a node"


def test_the_budget_is_per_address_not_global(pairing_world) -> None:
    """One noisy machine must not lock out every other node on the network."""

    repository, organization = pairing_world
    noisy = f"10.2.0.{uuid.uuid4().int % 250 + 1}"
    quiet = f"10.3.0.{uuid.uuid4().int % 250 + 1}"

    for _ in range(PAIRING_ATTEMPT_BUDGET + 5):
        repository.redeem_pairing_code("000000", "x", source_address=noisy)

    assert repository.redeem_pairing_code("000000", "x", source_address=quiet)["code"] != (
        "pairing_rate_limited"
    )


def test_a_rejected_attempt_leaves_no_node_behind(pairing_world) -> None:
    """The last line: a refused pairing must leave the system exactly as it was."""

    repository, organization = pairing_world
    name = f"clean-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    repository.redeem_pairing_code(minted["code"], "wrong-name")
    repository.redeem_pairing_code("000000", name)

    with repository._connection() as connection:
        nodes = connection.execute(
            "SELECT count(*) FROM public.compute_nodes WHERE name = %s", (name,)
        ).fetchone()["count"]
        credentials = connection.execute(
            "SELECT count(*) FROM public.node_credentials c JOIN public.compute_nodes n "
            "ON n.id = c.compute_node_id WHERE n.name = %s",
            (name,),
        ).fetchone()["count"]
    assert nodes == 0
    assert credentials == 0


# -- bounds on the code's own lifetime ---------------------------------------


def test_the_ttl_is_bounded_at_both_ends(pairing_world) -> None:
    """Long enough to walk to a machine, short enough not to be a credential for a day."""

    import pytest as _pytest

    from umbratrans_web.errors import InputError

    repository, organization = pairing_world
    for bad in (30, 3600):
        with _pytest.raises(InputError):
            _mint(repository, organization, f"ttl-{uuid.uuid4().hex[:8]}", ttl_seconds=bad)


def test_the_default_ttl_is_ten_minutes(pairing_world) -> None:
    """Asserted because the document and the code have to agree."""

    repository, organization = pairing_world
    minted = _mint(repository, organization, f"default-{uuid.uuid4().hex[:8]}")
    remaining = minted["expires_at"] - (_now_for_test())
    assert timedelta(seconds=550) < remaining <= timedelta(seconds=610), remaining


def _now_for_test():
    from umbratrans_web.compute import _now

    return _now()


def test_a_concurrent_redemption_produces_exactly_one_node(pairing_world) -> None:
    """Single use under a race, not just twice in sequence.

    There are two guards on this: the candidate query skips spent codes, and the claim
    is `UPDATE ... WHERE used_at IS NULL`. Removing the second leaves the first, so a
    sequential "redeem twice" test still passes - which is exactly what happened when
    this file was written. Sequential redemption cannot tell the two apart; only a
    concurrent one can, because it is the case where the candidate list was read before
    either claim landed.

    Two threads, one code, a barrier so both read the candidate list before either
    claims. Exactly one must succeed and exactly one node must exist.
    """

    import threading

    repository, organization = pairing_world
    name = f"race-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)

    barrier = threading.Barrier(2)
    results: list[dict] = []
    lock = threading.Lock()

    def attempt() -> None:
        # A separate repository per thread: sharing one pool across threads is its own
        # source of confusion, and the race being tested is in the database, not the pool.
        thread_repository = ComputeRepository(_SCRATCH, min_size=1, max_size=2, passwords=None)
        try:
            barrier.wait(timeout=10)
            outcome = thread_repository.redeem_pairing_code(minted["code"], name)
            with lock:
                results.append(outcome)
        finally:
            thread_repository.close()

    threads = [threading.Thread(target=attempt) for _ in range(2)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=30)

    assert len(results) == 2, [r.get("code") for r in results]
    successes = [r for r in results if r["ok"]]
    assert len(successes) == 1, f"{len(successes)} redemptions succeeded from one code"

    with repository._connection() as connection:
        count = connection.execute(
            "SELECT count(*) FROM public.compute_nodes WHERE name = %s", (name,)
        ).fetchone()["count"]
    assert count == 1, f"a raced code produced {count} nodes"


def test_a_spent_code_is_not_offered_as_a_candidate(pairing_world) -> None:
    """The filter that made the sequential single-use test unable to fail.

    Worth its own assertion because it is the *first* of the two guards, and the one
    that does the work in every ordinary case. If this were dropped, redemption of a
    spent code would reach the claim, and only the claim's `used_at IS NULL` would stop
    it - one mechanism instead of two, with the race window opened up accordingly.
    """

    repository, organization = pairing_world
    name = f"filter-{uuid.uuid4().hex[:8]}"
    minted = _mint(repository, organization, name)
    repository.redeem_pairing_code(minted["code"], name)

    with repository._connection() as connection:
        live = connection.execute(
            "SELECT count(*) FROM public.node_pairing_codes WHERE used_at IS NULL AND id = %s",
            (minted["id"],),
        ).fetchone()["count"]
    assert live == 0, "a spent code is still offered as a candidate"