plm-lottery/app/electrum/client.py

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


logger = logging.getLogger(__name__)


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

# Longest reply line the client accepts (B-75, ELC-20). asyncio's default StreamReader
# limit is 64 KiB, and a longer line used to raise out of the read loop and close the
# only connection: the pool address's get_history passed that after ~800 transactions
# (and a few hundred dust outputs sent to any public address do it on purpose), after
# which every confirmation poll dropped the connection and payouts never confirmed.
_MAX_LINE_BYTES = 16 * 1024 * 1024
_READ_CHUNK_BYTES = 64 * 1024
# How much of each end of a discarded line is kept to find the reply's `id` in, so the
# request it answered can fail at once instead of waiting out the 15s timeout.
_ID_SNIFF_BYTES = 256
_ID_PATTERN = re.compile(rb'"id"\s*:\s*(\d+)')


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, limit=_MAX_LINE_BYTES
        )
        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 get_history(self, scripthash: str) -> list[dict]:
        """Every transaction touching `scripthash`, each as {"tx_hash", "height"} —
        height > 0 means confirmed at that height, height <= 0 means still in the
        mempool. Used instead of blockchain.transaction.get's verbose=True mode
        for confirmation/existence checks (B-41): several Electrum server
        implementations and versions reject the verbose flag outright ("verbose
        transactions are currently unsupported"), while get_history is a plain,
        universally-supported method every server must implement.
        """
        return await self.request("blockchain.scripthash.get_history", [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:
        """Frame the stream into lines by hand rather than with StreamReader.readline:
        readline raises on a line over the limit, and that exception used to end this
        loop and with it the connection. An oversized line is now dropped up to its
        newline (B-75, ELC-20) and the connection stays up."""
        assert self._reader is not None
        buffer = bytearray()
        discarding = False
        head = b""
        tail = b""
        try:
            while True:
                chunk = await self._reader.read(_READ_CHUNK_BYTES)
                if not chunk:
                    break
                while chunk:
                    newline = chunk.find(b"\n")
                    piece = chunk if newline < 0 else chunk[:newline]
                    chunk = b"" if newline < 0 else chunk[newline + 1 :]
                    if discarding:
                        tail = (tail + piece)[-_ID_SNIFF_BYTES:]
                    else:
                        buffer += piece
                        if len(buffer) > _MAX_LINE_BYTES:
                            discarding = True
                            head = bytes(buffer[:_ID_SNIFF_BYTES])
                            tail = bytes(buffer[-_ID_SNIFF_BYTES:])
                            buffer.clear()
                    if newline >= 0:
                        if discarding:
                            self._fail_unreadable_reply(head + b" " + tail, "reply exceeds the line limit")
                            discarding = False
                        else:
                            self._handle_line(bytes(buffer))
                        buffer.clear()
            if buffer and not discarding:
                self._handle_line(bytes(buffer))  # a final line the server never terminated
        finally:
            # Whatever ended this loop — clean EOF, 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 _handle_line(self, line: bytes) -> None:
        if not line.strip():
            return
        try:
            message = json.loads(line)
        except ValueError:
            # A reply that is not JSON is a problem of that one request, not of the
            # connection (B-75).
            self._fail_unreadable_reply(line[:_ID_SNIFF_BYTES] + b" " + line[-_ID_SNIFF_BYTES:], "reply is not valid JSON")
            return
        if not isinstance(message, dict):
            logger.warning("ignoring an Electrum message that is not a JSON object")
            return
        self._dispatch(message)

    def _fail_unreadable_reply(self, sample: bytes, reason: str) -> None:
        """Fail the request an unusable reply belongs to, when its `id` can be read
        from `sample` (the first and last bytes of the line); otherwise that request
        fails by its own timeout (ELC-04)."""
        logger.warning("dropping an unusable Electrum reply: %s", reason)
        for match in _ID_PATTERN.finditer(sample):
            future = self._pending.pop(int(match.group(1)), None)
            if future is not None:
                if not future.done():
                    future.set_exception(ElectrumError(reason))
                return

    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"))