plm-lottery/app/electrum/client.py

import asyncio
import itertools
import json
import ssl
from typing import NamedTuple


class ElectrumEndpoint(NamedTuple):
    """One server to connect to. The listener rotates over a list of these so a
    single unreachable or misbehaving server doesn't take the platform down —
    every consumer of PLM chain data goes through this one connection."""

    host: str
    port: int
    use_ssl: bool = True

    def __str__(self) -> str:
        return f"{self.host}:{self.port}{'' if self.use_ssl else ' (plaintext)'}"


def parse_endpoints(
    primary_host: str, primary_port: int, primary_use_ssl: bool, fallback_spec: str
) -> list[ElectrumEndpoint]:
    """Build the connection rotation: the primary server first, then whatever
    ELECTRUM_FALLBACK_SERVERS lists.

    `fallback_spec` is comma-separated, each entry `host:port` (TLS, the normal
    case) or `host:port:notls`. Malformed entries raise ValueError rather than
    being skipped: a typo in a fallback server is something to fix at startup,
    not to discover during an outage, when the fallback is what's needed.
    Duplicates are dropped, keeping first position.
    """
    endpoints = [ElectrumEndpoint(primary_host, primary_port, primary_use_ssl)]
    for raw in fallback_spec.split(","):
        entry = raw.strip()
        if not entry:
            continue
        parts = entry.split(":")
        if len(parts) not in (2, 3):
            raise ValueError(f"invalid Electrum server {entry!r}: expected host:port[:notls]")
        host, port = parts[0].strip(), parts[1].strip()
        if not host or not port.isdigit():
            raise ValueError(f"invalid Electrum server {entry!r}: expected host:port[:notls]")
        use_ssl = True
        if len(parts) == 3:
            flag = parts[2].strip().lower()
            if flag not in ("ssl", "tls", "notls", "plain"):
                raise ValueError(f"invalid TLS flag {flag!r} in Electrum server {entry!r}")
            use_ssl = flag in ("ssl", "tls")
        endpoints.append(ElectrumEndpoint(host, int(port), use_ssl))

    deduped: list[ElectrumEndpoint] = []
    for endpoint in endpoints:
        if endpoint not in deduped:
            deduped.append(endpoint)
    return deduped



# Every request is bounded: without this, a half-open socket (the peer vanished
# without a FIN, or the read loop died — see _read_loop) leaves `await future`
# hanging forever, and with it whatever was awaiting the call. That used to be
# unbounded, which meant a POST /bets could hang while holding the per-user lock
# and the confirmation poller could stop polling permanently.
_REQUEST_TIMEOUT_SECONDS = 15


class ElectrumError(Exception):
    pass


class ElectrumClient:
    """Minimal asyncio Electrum protocol client: line-delimited JSON-RPC over TLS.

    Push notifications (blockchain.headers.subscribe, blockchain.scripthash.subscribe)
    arrive under the *same* method name as the subscribe call, multiplexed for every
    scripthash subscribed — callers read `notifications(method)` and, for scripthash
    pushes, dispatch on `params[0]` (the scripthash) themselves.

    A dead connection is observable rather than silent: `wait_closed()` resolves as
    soon as the read loop terminates for any reason, which is what lets
    ElectrumListener notice the drop and reconnect instead of waiting forever on
    notification queues nobody will ever fill again.
    """

    def __init__(self, host: str, port: int, use_ssl: bool = True):
        self.host = host
        self.port = port
        self.use_ssl = use_ssl
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._id_counter = itertools.count(1)
        self._pending: dict[int, asyncio.Future] = {}
        self._subscriptions: dict[str, asyncio.Queue] = {}
        self._read_task: asyncio.Task | None = None
        self._closed = asyncio.Event()

    async def connect(self) -> None:
        # Electrum servers commonly present self-signed certs; the protocol's trust
        # model is server consensus, not TLS PKI, so we only use SSL for transport
        # encryption and don't verify the certificate chain/hostname.
        ssl_context = None
        if self.use_ssl:
            ssl_context = ssl.create_default_context()
            ssl_context.check_hostname = False
            ssl_context.verify_mode = ssl.CERT_NONE
        self._reader, self._writer = await asyncio.open_connection(self.host, self.port, ssl=ssl_context)
        self._read_task = asyncio.create_task(self._read_loop())
        await self.request("server.version", ["plm-lottery", "1.4"])

    async def close(self) -> None:
        self._closed.set()
        if self._read_task is not None:
            self._read_task.cancel()
        if self._writer is not None:
            self._writer.close()
            try:
                await asyncio.wait_for(self._writer.wait_closed(), timeout=2)
            except (ssl.SSLError, TimeoutError, asyncio.TimeoutError, ConnectionError):
                pass  # some Electrum servers don't send a clean TLS close_notify

    async def wait_closed(self) -> None:
        """Resolves once this connection is gone — read loop finished (peer closed,
        protocol error, TLS failure) or close() was called. ElectrumListener races
        this against its notification consumers so a drop triggers a reconnect."""
        await self._closed.wait()

    async def request(self, method: str, params: list | None = None) -> object:
        if self._writer is None or self._closed.is_set():
            raise ElectrumError("not connected")
        request_id = next(self._id_counter)
        future: asyncio.Future = asyncio.get_running_loop().create_future()
        self._pending[request_id] = future
        payload = json.dumps({"id": request_id, "method": method, "params": params or []}) + "\n"
        try:
            self._writer.write(payload.encode())
            await self._writer.drain()
        except (ConnectionError, ssl.SSLError, OSError) as exc:
            self._pending.pop(request_id, None)
            self._closed.set()
            raise ElectrumError(f"write failed: {exc}") from exc
        try:
            return await asyncio.wait_for(future, timeout=_REQUEST_TIMEOUT_SECONDS)
        except (TimeoutError, asyncio.TimeoutError) as exc:
            self._pending.pop(request_id, None)
            # A server that owes us a reply and never sends one is indistinguishable
            # from a dead socket, and retrying on the same connection would keep
            # hitting it — tear it down so the listener reconnects.
            self._closed.set()
            raise ElectrumError(f"{method} timed out after {_REQUEST_TIMEOUT_SECONDS}s") from exc

    async def ping(self) -> None:
        """Keepalive: Electrum servers drop idle connections (commonly after ~10
        minutes), which on a quiet instance would otherwise be the normal way this
        connection dies. Called periodically by ElectrumListener."""
        await self.request("server.ping")

    def notifications(self, method: str) -> asyncio.Queue:
        return self._subscriptions.setdefault(method, asyncio.Queue())

    async def subscribe_headers(self) -> dict:
        self.notifications("blockchain.headers.subscribe")
        return await self.request("blockchain.headers.subscribe")

    async def subscribe_scripthash(self, scripthash: str) -> str | None:
        self.notifications("blockchain.scripthash.subscribe")
        return await self.request("blockchain.scripthash.subscribe", [scripthash])

    async def listunspent(self, scripthash: str) -> list[dict]:
        return await self.request("blockchain.scripthash.listunspent", [scripthash])

    async def broadcast(self, raw_tx_hex: str) -> str:
        return await self.request("blockchain.transaction.broadcast", [raw_tx_hex])

    async def get_transaction(self, txid: str, verbose: bool = False) -> object:
        return await self.request("blockchain.transaction.get", [txid, verbose])

    async def _read_loop(self) -> None:
        assert self._reader is not None
        try:
            while True:
                line = await self._reader.readline()
                if not line:
                    break
                message = json.loads(line)
                self._dispatch(message)
        finally:
            # Whatever ended this loop — clean EOF, protocol error, TLS failure — the
            # connection is unusable from here on. Setting this is what makes the
            # death observable to wait_closed(), and so to the listener's reconnect
            # logic; without it the listener waited on notification queues nobody
            # would ever fill again, forever (B-01).
            self._closed.set()
            error = ElectrumError("connection closed")
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(error)
            self._pending.clear()

    def _dispatch(self, message: dict) -> None:
        message_id = message.get("id")
        if message_id is not None and message_id in self._pending:
            future = self._pending.pop(message_id)
            if future.done():
                return
            if message.get("error"):
                future.set_exception(ElectrumError(message["error"]))
            else:
                future.set_result(message.get("result"))
        elif "method" in message:
            queue = self._subscriptions.get(message["method"])
            if queue is not None:
                queue.put_nowait(message.get("params"))