interface: hardening against resource exhaustion (take1)

electrum/interface.py

@@ -126,9 +126,13 @@ def assert_hex_str(val: Any) -> None:
         raise RequestCorrupted(f'{val!r} should be a hex str')
 
 
-def assert_dict_contains_field(d: Any, *, field_name: str) -> Any:
+def assert_dict(d: Any) -> None:
     if not isinstance(d, dict):
         raise RequestCorrupted(f'{d!r} should be a dict')
+
+
+def assert_dict_contains_field(d: Any, *, field_name: str) -> Any:
+    assert_dict(d)
     if field_name not in d:
         raise RequestCorrupted(f'required field {field_name!r} missing from dict')
     return d[field_name]
@@ -161,36 +165,75 @@ class ChainResolutionMode(enum.Enum):
 
 class NotificationSession(RPCSession):
 
+    INC_REQ_CONCURRENCY_FOR_MAIN_IFACE = 5
+    INC_REQ_CONCURRENCY_FOR_OTHER_IFACE = 2
+
     def __init__(self, *args, interface: 'Interface', **kwargs):
         super(NotificationSession, self).__init__(*args, **kwargs)
-        self.subscriptions = defaultdict(list)
-        self.cache = {}
+        self.subscriptions = defaultdict(list)  # type: defaultdict[str, list[asyncio.Queue]]
+        self.subs_cache = {}  # type: dict[str, Any]
         self._msg_counter = itertools.count(start=1)
         self.interface = interface
         self.taskgroup = interface.taskgroup
-        self.cost_hard_limit = 0  # disable aiorpcx resource limits
+        assert hasattr(self, "cost_hard_limit")  # in base class
+        self.cost_hard_limit = 0  # disable aiorpcx resource limits  # TODO secondaries
+        assert hasattr(self, "initial_concurrent")  # in base class
+        self.initial_concurrent = self.INC_REQ_CONCURRENCY_FOR_OTHER_IFACE
+        self._incoming_concurrency.set_target(self.INC_REQ_CONCURRENCY_FOR_OTHER_IFACE)  # this limits memory usage
 
         # To log pre-processed json traffic, uncomment:
         #self.logger.setLevel(logging.DEBUG)  # from aiorpcx
         #self.verbosity = 4
 
     async def handle_request(self, request):
         self.maybe_log(f"--> {request}")
+        # note: the caller enforces a timeout (processing_timeout) on us, after which we will get cancelled
+        # note: size of a single request is bounded by the framer: config.NETWORK_MAX_INCOMING_MSG_SIZE.
+        #       That bound is too relaxed, to allow requesting a large tx. Most messages should
+        #       be MUCH smaller, especially notifications...
         try:
             if isinstance(request, Notification):
                 params, result = request.args[:-1], request.args[-1]
                 key = self.get_hashable_key_for_rpc_call(request.method, params)
                 if key in self.subscriptions:
-                    self.cache[key] = result
+                    self.subs_cache[key] = result
                     for queue in self.subscriptions[key]:
+                        # note: if queue is full, queue.put will block until a free slot is available.
+                        #       This limits memory usage.
+                        # note: there are up to "initial_concurrent" handle_request() calls blocking here.
                         await queue.put(request.args)
                 else:
                     raise Exception(f'unexpected notification')
             else:
                 raise Exception(f'unexpected request. not a notification')
-        except Exception as e:
+        except BaseException as e:
+            # note: we must not silently drop any notification, so we even catch CancelledError(BaseException).
+            #       Instead, make sure we disconnect.
             self.interface.logger.info(f"error handling request {request}. exc: {repr(e)}")
             await self.close()
+            if isinstance(e, asyncio.CancelledError):
+                raise
+
+    async def _throttled_request(self, request):
+        # If we are already processing too many received messages, pause reading from the transport.
+        # This limits memory usage.
+        # FIXME this is a horrible ugly hack, accessing deep internals inside aiorpcx.
+        #       A proper fix instead should be implemented inside aiorpcx.
+        #       A proper fix should also be looking at the cumulative byte size of pending requests,
+        #       instead of request count.
+        # note: we get called for incoming Requests and Notifications. (not for Responses)
+        #       As we are not server, we don't expect getting Requests: handle_request will raise and close().
+        rstransport = self.transport  # type: PaddedRSTransport
+        asyncio_transport = rstransport._asyncio_transport  # type: asyncio.Transport
+        if self._incoming_concurrency._semaphore.locked():
+            asyncio_transport.pause_reading()
+        try:
+            return await super()._throttled_request(request)
+        finally:
+            # Just finished processing one received request. If there are not too many unprocessed ones
+            # AND the outgoing send buffer has room, we can resume reading:
+            if not self._incoming_concurrency._semaphore.locked() and rstransport._can_send.is_set():
+                asyncio_transport.resume_reading()
 
     async def send_request(self, *args, timeout=None, **kwargs):
         # note: semaphores/timeouts/backpressure etc are handled by
@@ -223,15 +266,23 @@ def set_default_timeout(self, timeout):
         self.max_send_delay = timeout
 
     async def subscribe(self, method: str, params: List, queue: asyncio.Queue):
-        # note: until the cache is written for the first time,
-        # each 'subscribe' call might make a request on the network.
+        assert queue.maxsize > 0, "infinite queue maxsize not allowed"
         key = self.get_hashable_key_for_rpc_call(method, params)
-        self.subscriptions[key].append(queue)
-        if key in self.cache:
-            result = self.cache[key]
+        if key in self.subs_cache:
+            result = self.subs_cache[key]
         else:
+            # note: until subs_cache is written for the first time,
+            #       each 'subscribe' call might make a request on the network.
             result = await self.send_request(method, params)
-            self.cache[key] = result
+            self.subs_cache[key] = result
+        # note: multiple Synchronizers (from different Wallet objects) might sub to the same key,
+        #       hence subscriptions map key->list[queue]
+        # note: only after we got the initial response from the network, we save the subscription.
+        #       This way we disallow force all notifications to arrive after the initial response.
+        #       (as the queue size is bounded, that could even cause a deadlock)
+        self.subscriptions[key].append(queue)
+        # note: if queue is full, queue.put will block until a free slot is available.
+        #       This limits memory usage.  FIXME add timeout??
         await queue.put(params + [result])
 
     def unsubscribe(self, queue):
@@ -584,6 +635,7 @@ def __init__(self, *, network: 'Network', server: ServerAddr):
         # Failing verification will get the interface closed.
         self.tip_header = None  # type: Optional[dict]
         self.tip = 0
+        self._tip_unprocessed_evt = asyncio.Event()
 
         self._headers_cache = {}  # type: Dict[int, bytes]
         self._rawtx_cache = LRUCache(maxsize=20)  # type: LRUCache[str, bytes]  # txid->rawtx
@@ -976,6 +1028,12 @@ def is_main_server(self) -> bool:
         return (self.network.interface == self or
                 self.network.interface is None and self.network.default_server == self.server)
 
+    def mark_as_main_server(self) -> None:
+        """Called when the network switches to this interface."""
+        assert self.session
+        self.session.initial_concurrent = self.session.INC_REQ_CONCURRENCY_FOR_MAIN_IFACE
+        self.session._incoming_concurrency.set_target(self.session.INC_REQ_CONCURRENCY_FOR_MAIN_IFACE)
+
     async def open_session(
         self,
         *,
@@ -1023,7 +1081,8 @@ async def open_session(
                 async with self.taskgroup as group:
                     await group.spawn(self.ping)
                     await group.spawn(self.request_fee_estimates)
-                    await group.spawn(self.run_fetch_blocks)
+                    await group.spawn(self._subscribe_to_headers)
+                    await group.spawn(self._loop_process_header_at_tip)
                     await group.spawn(self.monitor_connection)
             except aiorpcx.jsonrpc.RPCError as e:
                 if e.code in (
@@ -1088,30 +1147,43 @@ async def close(self, *, force_after: int = None):
             await self.session.close(force_after=force_after)
         # monitor_connection will cancel tasks
 
-    async def run_fetch_blocks(self):
-        header_queue = asyncio.Queue()
-        await self.session.subscribe('blockchain.headers.subscribe', [], header_queue)
+    async def _subscribe_to_headers(self):
+        unsanitized_header_queue = asyncio.Queue(maxsize=1)  # maxsize limits memory usage
+        await self.session.subscribe('blockchain.headers.subscribe', [], unsanitized_header_queue)
         while True:
-            item = await header_queue.get()
-            raw_header = item[0]
-            height = raw_header['height']
-            header_bytes = bfh(raw_header['hex'])
+            item = await unsanitized_header_queue.get()
+            # parse response
+            assert len(item) == 1
+            resp_header = item[0]
+            assert_dict(resp_header)
+            height = assert_dict_contains_field(resp_header, field_name='height')
+            assert_non_negative_integer(height)
+            header_hex = assert_dict_contains_field(resp_header, field_name='hex')
+            header_bytes = bfh(header_hex)
+            assert len(resp_header) == 2, f"resp_header contains redundant fields. got {resp_header.keys()}"
             header_dict = blockchain.deserialize_header(header_bytes, height)
+            # process header
             self.tip_header = header_dict
             self.tip = height
             if self.tip < constants.net.max_checkpoint():
                 raise GracefulDisconnect(
                     f"server tip below max checkpoint. ({self.tip} < {constants.net.max_checkpoint()})")
-            self._mark_ready()
             self._headers_cache.clear()  # tip changed, so assume anything could have happened with chain
             self._headers_cache[height] = header_bytes
+            self._tip_unprocessed_evt.set()
+
+    async def _loop_process_header_at_tip(self):
+        while True:
+            # wait until tip changes
+            await self._tip_unprocessed_evt.wait()
+            self._mark_ready()
             try:
                 blockchain_updated = await self._process_header_at_tip()
             finally:
                 self._headers_cache.clear()  # to reduce memory usage
             # header processing done
             if self.is_main_server() or blockchain_updated:
-                self.logger.info(f"new chain tip. {height=}")
+                self.logger.info(f"new chain tip. height={self.tip}")
             if blockchain_updated:
                 util.trigger_callback('blockchain_updated')
                 self._blockchain_updated.set()
@@ -1127,7 +1199,7 @@ async def _process_header_at_tip(self) -> bool:
         True - new header we didn't have, or reorg
         """
         height, header = self.tip, self.tip_header
-        async with self.network.bhi_lock:
+        async with self.network.bhi_lock:  # FIXME secondary server can starve main
             if self.blockchain.height() >= height and self.blockchain.check_header(header):
                 # another interface amended the blockchain
                 return False

electrum/network.py

@@ -891,6 +891,7 @@ async def switch_to_interface(self, server: ServerAddr):
 
         # Stop any current interface in order to terminate subscriptions,
         # and to cancel tasks in interface.taskgroup.
+        # This also indirectly undoes i.mark_as_main_server().
         if old_server and old_server != server:
             # don't wait for old_interface to close as that might be slow:
             await self.taskgroup.spawn(self._close_interface(old_interface))
@@ -907,6 +908,7 @@ async def switch_to_interface(self, server: ServerAddr):
             self.logger.info(f"switching to {server}")
             blockchain_updated = i.blockchain != self.blockchain()
             self.interface = i
+            i.mark_as_main_server()
             try:
                 await i.taskgroup.spawn(self._request_server_info(i))
             except RuntimeError as e:  # see #7677

electrum/scripts/block_headers.py

@@ -21,7 +21,7 @@
     time.sleep(1)
     print_msg("waiting for network to get connected...")
 
-header_queue = asyncio.Queue()
+header_queue = asyncio.Queue(maxsize=1)
 
 @log_exceptions
 async def f():

electrum/synchronizer.py

@@ -35,7 +35,7 @@
 from .util import make_aiohttp_session, NetworkJobOnDefaultServer, random_shuffled_copy, OldTaskGroup
 from .bitcoin import address_to_scripthash, is_address, neuter_bitcoin_address
 from .logging import Logger
-from .interface import GracefulDisconnect, NetworkTimeout
+from .interface import GracefulDisconnect, NetworkTimeout, assert_hash256_str
 
 if TYPE_CHECKING:
     from .network import Network
@@ -71,7 +71,7 @@ def _reset(self):
         self.scripthash_to_address = {}
         self._processed_some_notifications = False  # so that we don't miss them
         # Queues
-        self.status_queue = asyncio.Queue()
+        self.status_queue = asyncio.Queue(maxsize=15)  # maxsize limits memory usage
 
     async def _run_tasks(self, *, taskgroup):
         await super()._run_tasks(taskgroup=taskgroup)
@@ -117,8 +117,13 @@ async def _subscribe_to_address(self, addr):
 
     async def handle_status(self):
         while True:
-            h, status = await self.status_queue.get()
-            addr = self.scripthash_to_address[h]
+            sh, status = await self.status_queue.get()
+            # basic checks for response
+            assert_hash256_str(sh)
+            if status is not None:
+                assert_hash256_str(status)
+            # process status
+            addr = self.scripthash_to_address[sh]
             self._handling_addr_statuses.add(addr)
             self.requested_addrs.discard(addr)  # ok for addr not to be present
             await self.taskgroup.spawn(self._on_address_status, addr, status)