plm-lottery/tests/unit/test_electrum_client.py

import asyncio
import json

from app.electrum.client import ElectrumClient, ElectrumError


class FakeWriter:
    def __init__(self):
        self.written = b""

    def write(self, data: bytes) -> None:
        self.written += data

    async def drain(self) -> None:
        pass

    def close(self) -> None:
        pass

    async def wait_closed(self) -> None:
        pass


async def _client_with_fake_transport() -> tuple[ElectrumClient, asyncio.StreamReader, FakeWriter]:
    client = ElectrumClient("localhost", 1234)
    reader = asyncio.StreamReader()
    writer = FakeWriter()
    client._reader = reader
    client._writer = writer
    client._read_task = asyncio.create_task(client._read_loop())
    return client, reader, writer


async def test_request_resolves_on_matching_response():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("blockchain.headers.subscribe"))
    await asyncio.sleep(0)  # let request() write the payload
    sent = json.loads(writer.written.decode())
    assert sent["method"] == "blockchain.headers.subscribe"

    reader.feed_data((json.dumps({"id": sent["id"], "result": {"height": 100}}) + "\n").encode())
    result = await task
    assert result == {"height": 100}

    client._read_task.cancel()


async def test_error_response_raises_electrum_error():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("blockchain.transaction.broadcast", ["deadbeef"]))
    await asyncio.sleep(0)
    sent = json.loads(writer.written.decode())

    reader.feed_data((json.dumps({"id": sent["id"], "error": "bad tx"}) + "\n").encode())
    try:
        await task
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass

    client._read_task.cancel()


async def test_notification_delivered_to_subscription_queue():
    client, reader, writer = await _client_with_fake_transport()
    queue = client.notifications("blockchain.scripthash.subscribe")

    push = {"method": "blockchain.scripthash.subscribe", "params": ["abcd", "newstatus"]}
    reader.feed_data((json.dumps(push) + "\n").encode())
    await asyncio.sleep(0)

    params = await asyncio.wait_for(queue.get(), timeout=1)
    assert params == ["abcd", "newstatus"]

    client._read_task.cancel()


async def test_request_times_out_instead_of_hanging_forever(monkeypatch):
    """B-01: a server that owes us a reply and never sends one used to hang the
    caller permanently — which meant a POST /bets could hang while holding the
    per-user lock, and the confirmation poller could stop polling for good."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_REQUEST_TIMEOUT_SECONDS", 0.05)
    client, reader, writer = await _client_with_fake_transport()

    try:
        await client.request("blockchain.transaction.get", ["deadbeef"])
        assert False, "expected ElectrumError"
    except ElectrumError as exc:
        assert "timed out" in str(exc)

    # The connection is torn down, so callers stop reusing a server that owes us.
    assert client._closed.is_set()
    client._read_task.cancel()


async def test_wait_closed_resolves_when_the_read_loop_dies():
    """B-01: the read loop dying used to be invisible — the listener sat on its
    notification queues forever and never reconnected."""
    client, reader, writer = await _client_with_fake_transport()

    waiter = asyncio.create_task(client.wait_closed())
    reader.feed_eof()  # peer closed the connection

    await asyncio.wait_for(waiter, timeout=1)
    client._read_task.cancel()


async def test_pending_request_fails_when_the_connection_drops():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("server.ping"))
    await asyncio.sleep(0)
    reader.feed_eof()

    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass


def test_parse_endpoints_puts_the_primary_first_and_dedupes():
    from app.electrum.client import ElectrumEndpoint, parse_endpoints

    endpoints = parse_endpoints(
        "primary.example", 50002, True, "second.example:50002, third.example:50001:notls, primary.example:50002"
    )
    assert endpoints == [
        ElectrumEndpoint("primary.example", 50002, True),
        ElectrumEndpoint("second.example", 50002, True),
        ElectrumEndpoint("third.example", 50001, False),
    ]


def test_parse_endpoints_handles_an_empty_fallback_list():
    from app.electrum.client import ElectrumEndpoint, parse_endpoints

    assert parse_endpoints("only.example", 50002, True, "") == [ElectrumEndpoint("only.example", 50002, True)]
    assert parse_endpoints("only.example", 50002, True, "  ,  ") == [
        ElectrumEndpoint("only.example", 50002, True)
    ]


def test_parse_endpoints_rejects_a_malformed_entry():
    """A typo in a fallback server must fail at startup, not during the outage when
    the fallback is the thing that's needed."""
    import pytest

    from app.electrum.client import parse_endpoints

    for bad in ["nohost:", "host:notaport", "host", "host:50002:weird"]:
        with pytest.raises(ValueError):
            parse_endpoints("primary.example", 50002, True, bad)


# --- B-75 (ELC-20): an oversized or malformed reply is a request error, never a
# dead connection. A tiny line limit stands in for 16 MiB. ---------------------


async def _send_request(client, writer, reader):
    task = asyncio.create_task(client.request("blockchain.scripthash.get_history", ["ab"]))
    await asyncio.sleep(0)
    return task, json.loads(writer.written.decode().splitlines()[-1])["id"]


async def test_the_line_limit_is_16_mib_and_connect_applies_it(monkeypatch):
    """ELC-20 / B-75: the StreamReader is opened with a 16 MiB limit, not 64 KiB."""
    from app.electrum import client as client_module

    assert client_module._MAX_LINE_BYTES == 16 * 1024 * 1024
    seen = {}

    async def fake_open_connection(host, port, **kwargs):
        seen.update(kwargs)
        raise OSError("stop here")

    monkeypatch.setattr(asyncio, "open_connection", fake_open_connection)
    try:
        await ElectrumClient("localhost", 1).connect()
    except OSError:
        pass
    assert seen["limit"] == 16 * 1024 * 1024


async def test_an_oversized_reply_fails_its_request_and_keeps_the_connection(monkeypatch):
    """ELC-20 / B-75: the reply is discarded to its newline (even when it arrives in
    pieces), the request whose id is in its head fails at once, and the next reply
    on the same connection is still delivered."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_MAX_LINE_BYTES", 1024)
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    head = b'{"id": %d, "result": [' % request_id
    reader.feed_data(head + b"1," * 600)  # already over the limit, no newline yet
    await asyncio.sleep(0)
    reader.feed_data(b"1," * 600 + b"1]}\n")
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    task2, request_id2 = await _send_request(client, writer, reader)
    reader.feed_data((json.dumps({"id": request_id2, "result": "ok"}) + "\n").encode())
    assert await asyncio.wait_for(task2, timeout=1) == "ok"

    client._read_task.cancel()


async def test_an_oversized_reply_with_the_id_at_the_end_also_fails_its_request(monkeypatch):
    """ELC-20: the id is looked for in the last bytes of the line too."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_MAX_LINE_BYTES", 1024)
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    reader.feed_data(b'{"result": [' + b"1," * 1000 + b'1], "id": %d}\n' % request_id)
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    client._read_task.cancel()


async def test_a_reply_that_is_not_json_fails_its_request_not_the_connection():
    """ELC-20 / B-75: json.loads used to raise out of the read loop and close it."""
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    reader.feed_data(b'{"id": %d, "result": [1, 2,\n' % request_id)
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    # Garbage with no readable id, and a non-object message, are dropped silently.
    reader.feed_data(b"not json at all\n[1, 2]\n")
    task2, request_id2 = await _send_request(client, writer, reader)
    reader.feed_data((json.dumps({"id": request_id2, "result": 7}) + "\n").encode())
    assert await asyncio.wait_for(task2, timeout=1) == 7

    client._read_task.cancel()


# --- B-84 (ELC-01, ELC-03): optional certificate pinning --------------------------

_PIN = "ab" * 32


def test_parse_endpoints_reads_pins_on_fallbacks_and_the_primary():
    """ELC-01 / ELC-03 / B-84: `host:port:pin=<sha256>` and `host:port:<mode>:pin=<sha256>`,
    plus ELECTRUM_CERT_PIN for the primary; the pin is normalised to lower case."""
    from app.electrum.client import ElectrumEndpoint, parse_endpoints

    endpoints = parse_endpoints(
        "primary.example",
        50002,
        True,
        f"a.example:50002:pin={_PIN.upper()}, b.example:50002:tls:pin={_PIN}, c.example:50001:notls",
        primary_pin=_PIN,
    )
    assert endpoints == [
        ElectrumEndpoint("primary.example", 50002, True, _PIN),
        ElectrumEndpoint("a.example", 50002, True, _PIN),
        ElectrumEndpoint("b.example", 50002, True, _PIN),
        ElectrumEndpoint("c.example", 50001, False, None),
    ]


def test_parse_endpoints_rejects_malformed_pins():
    """ELC-01 / B-84: a bad pin blocks the start, like any malformed entry."""
    import pytest

    from app.electrum.client import parse_endpoints

    for bad in [
        "h:1:pin=",
        "h:1:pin=abc",
        f"h:1:pin={'zz' * 32}",
        f"h:1:notls:pin={_PIN}",  # nothing to pin without TLS
        f"h:1:pin={_PIN}:pin={_PIN}",
        f"h:1:tls:ssl:pin={_PIN}",
        f"h:1:tls:ssl:notls:pin={_PIN}",  # too many fields
    ]:
        with pytest.raises(ValueError):
            parse_endpoints("primary.example", 50002, True, bad)
    with pytest.raises(ValueError):
        parse_endpoints("primary.example", 50002, True, "", primary_pin="nothex")
    with pytest.raises(ValueError):
        parse_endpoints("primary.example", 50002, False, "", primary_pin=_PIN)


class _FakeSslObject:
    def __init__(self, der):
        self._der = der

    def getpeercert(self, binary_form=False):
        assert binary_form
        return self._der


class _PinWriter(FakeWriter):
    def __init__(self, der):
        super().__init__()
        self._ssl = _FakeSslObject(der) if der is not None else None
        self.closed = False

    def get_extra_info(self, name):
        return self._ssl if name == "ssl_object" else None

    def close(self):
        self.closed = True


def _answer_features(reader, writer, result=None):
    """Make `writer` answer `server.features` the moment it is asked: a reply fed up front
    would be read before its request is registered, and dropped."""
    from app.wallet.plm_network import GENESIS_HASH

    original = writer.write

    def write(data):
        original(data)
        sent = json.loads(data.decode())
        if sent["method"] == "server.features":
            reply = {"genesis_hash": GENESIS_HASH} if result is None else result
            reader.feed_data((json.dumps({"id": sent["id"], "result": reply}) + "\n").encode())

    writer.write = write


async def _connect_with_pin(monkeypatch, pin, der):
    import hashlib  # noqa: F401

    writer = _PinWriter(der)
    reader = asyncio.StreamReader()
    reader.feed_data(b'{"id": 1, "result": ["fake", "1.4"]}\n')
    _answer_features(reader, writer)

    async def fake_open_connection(host, port, **kwargs):
        return reader, writer

    monkeypatch.setattr(asyncio, "open_connection", fake_open_connection)
    client = ElectrumClient("localhost", 1, True, pin)
    try:
        await client.connect()
    finally:
        if client._read_task is not None:
            client._read_task.cancel()
    return writer


async def test_a_connection_with_the_pinned_certificate_is_accepted(monkeypatch):
    """ELC-03 / B-84: the DER certificate hashes to the pin."""
    import hashlib

    der = b"the server certificate"
    writer = await _connect_with_pin(monkeypatch, hashlib.sha256(der).hexdigest().upper(), der)
    assert not writer.closed


async def test_a_connection_with_another_certificate_fails_like_an_unreachable_server(monkeypatch):
    """ELC-03 / B-84: mismatch (and no certificate at all) raise ElectrumError and close."""
    import hashlib

    import pytest

    pin = hashlib.sha256(b"the pinned certificate").hexdigest()
    for der in (b"someone else in the middle", None):
        writer = _PinWriter(der)

        async def fake_open_connection(host, port, writer=writer, **kwargs):
            return asyncio.StreamReader(), writer

        monkeypatch.setattr(asyncio, "open_connection", fake_open_connection)
        with pytest.raises(ElectrumError):
            await ElectrumClient("localhost", 1, True, pin).connect()
        assert writer.closed


async def test_main_builds_clients_with_the_endpoint_pin():
    """ELC-03 / B-84: app.main hands each endpoint's pin to its client."""
    from app.electrum.client import ElectrumEndpoint
    from app.main import _make_electrum_client

    assert _make_electrum_client(ElectrumEndpoint("h", 1, True, _PIN)).pin == _PIN
    assert _make_electrum_client(ElectrumEndpoint("h", 1, True)).pin is None


async def test_only_a_server_error_object_is_an_explicit_refusal():
    """B-76 / TX-04: a JSON-RPC `error` reply is an ElectrumRpcError carrying the
    server's wording (string or object); a timeout or a dropped connection is a plain
    ElectrumError — they say nothing about whether the server acted."""
    from app.electrum import client as client_module
    from app.electrum.client import ElectrumRpcError

    client, reader, writer = await _client_with_fake_transport()
    task = asyncio.create_task(client.request("blockchain.transaction.broadcast", ["00"]))
    await asyncio.sleep(0)
    sent = json.loads(writer.written.decode())
    reader.feed_data((json.dumps({"id": sent["id"], "error": {"code": -26, "message": "txn-mempool-conflict"}}) + "\n").encode())
    try:
        await task
        assert False, "expected ElectrumRpcError"
    except ElectrumRpcError as exc:
        assert exc.message == "txn-mempool-conflict"
    client._read_task.cancel()

    original = client_module._REQUEST_TIMEOUT_SECONDS
    client_module._REQUEST_TIMEOUT_SECONDS = 0.05
    try:
        silent, _reader, _writer = await _client_with_fake_transport()
        try:
            await silent.request("blockchain.transaction.broadcast", ["00"])
            assert False, "expected ElectrumError"
        except ElectrumError as exc:
            assert not isinstance(exc, ElectrumRpcError)
        silent._read_task.cancel()
    finally:
        client_module._REQUEST_TIMEOUT_SECONDS = original


async def test_a_server_on_another_chain_is_refused_at_the_handshake(monkeypatch):
    """ELC-30: Bitcoin's genesis (or no genesis at all) fails like an unreachable server."""
    import pytest

    bitcoin = "000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26e"
    for result in ({"genesis_hash": bitcoin}, {}):
        writer = FakeWriter()
        reader = asyncio.StreamReader()
        reader.feed_data(b'{"id": 1, "result": ["fake", "1.4"]}\n')
        _answer_features(reader, writer, result)

        async def fake_open_connection(host, port, reader=reader, writer=writer, **kwargs):
            return reader, writer

        monkeypatch.setattr(asyncio, "open_connection", fake_open_connection)
        with pytest.raises(ElectrumError, match="not on the PLM chain"):
            await ElectrumClient("localhost", 1).connect()