interface: hardening against resource exhaustion (take2)

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,27 +165,31 @@ class ChainResolutionMode(enum.Enum):
 
 class NotificationSession(RPCSession):
 
+    COST_INCOMING_REQUEST = 100
+
     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
+        self.set_strict_resource_limits()
 
         # To log pre-processed json traffic, uncomment:
         #self.logger.setLevel(logging.DEBUG)  # from aiorpcx
         #self.verbosity = 4
 
     async def handle_request(self, request):
+        # note: we get called for incoming Requests and Notifications. (not for Responses)
         self.maybe_log(f"--> {request}")
+        self.bump_cost(self.COST_INCOMING_REQUEST)
         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]:
                         await queue.put(request.args)
                 else:
@@ -223,15 +231,17 @@ 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.
         key = self.get_hashable_key_for_rpc_call(method, params)
+        # note: multiple Synchronizers (from different Wallet objects) might sub to the same key,
+        #       hence subscriptions map key->list[queue]
         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
         await queue.put(params + [result])
 
     def unsubscribe(self, queue):
@@ -271,6 +281,40 @@ async def close(self, *, force_after: int = None):
             force_after = 1  # seconds
         await super().close(force_after=force_after)
 
+    def set_strict_resource_limits(self):
+        # Apply strict resource limits to each interface.
+        # - This limits incoming bandwidth, and indirectly limits e.g. memory usage.
+        #   FIXME limit memory usage directly
+        #   - confusingly for our "client" use case, outgoing bandwidth is also counted,
+        #     however any meaningful limiting only happens in _throttled_message->_incoming_concurrency,
+        #     which only gets called on inc-notifications and inc-requests (reqs we should not receive though).
+        #   - processing an inc-notification or an inc-request also incurs COST_INCOMING_REQUEST.
+        # - These limits are intended for the non-main secondary interfaces,
+        #   as the main interface is expected to have a lot of traffic.
+        # - Secondary interfaces should generate minimal traffic:
+        #   as they are only used for polled fee estimates (minimal data) and header subs (minimal data).
+        #   Note a header notification can force us into fork resolution and downloading lots
+        #   of headers. The initial headers download might also happen on any interface.
+        #   Headers download uses non-trivial amounts of data, e.g. 100k x 160 hex headers take 16 MB.
+        assert hasattr(NotificationSession, "cost_hard_limit")  # in base class
+        self.bw_cost_per_byte = 1 / 1_000
+        self.cost_hard_limit = 30_000  # 30 MB of bandwidth, in+out
+        self.cost_soft_limit = self.cost_hard_limit - 1  # this effectively disables the soft limit
+        self.cost_decay_per_sec = self.cost_hard_limit / 600  # refund over 10 minutes
+
+    def remove_resource_limits(self):
+        assert hasattr(NotificationSession, "cost_hard_limit")  # in base class
+        # remove static limits:
+        self.cost_hard_limit = 0
+        self.cost_soft_limit = 0
+        # try to reset costs incurred so far (e.g. when promoting a non-main interface to main):
+        self._cost_last = self.cost = 0
+        self._cost_fraction = 0
+        self._incoming_concurrency.set_target(self.initial_concurrent)
+
+    def on_disconnect_due_to_excessive_session_cost(self):
+        self.interface.logger.info(f"closing session over resource usage. cost={self.cost}")
+
 
 class NetworkException(Exception): pass
 
@@ -976,6 +1020,11 @@ 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.remove_resource_limits()
+
     async def open_session(
         self,
         *,
@@ -1093,11 +1142,18 @@ async def run_fetch_blocks(self):
         await self.session.subscribe('blockchain.headers.subscribe', [], header_queue)
         while True:
             item = await header_queue.get()
-            raw_header = item[0]
-            height = raw_header['height']
-            header_bytes = bfh(raw_header['hex'])
+            # 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)
-            self.tip_header = header_dict
+            # process header
+            self.tip_header = header_dict  # TODO assert it got changed to something different?
             self.tip = height
             if self.tip < constants.net.max_checkpoint():
                 raise GracefulDisconnect(

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/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
@@ -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)

tests/test_interface.py

@@ -202,3 +202,39 @@ async def test_dont_request_gethistory_if_status_change_results_from_mempool_txs
             w1.adb.get_address_history(w1_addr),
             {funding_txid: server_blockheight})
 
+    async def test_we_disconnect_on_incoming_request(self):
+        """We don't expect the server to send us any requests of its own."""
+        interface = await self._start_iface_and_wait_for_sync()
+        with self.assertLogs('electrum', level='INFO') as logs:
+            with self.assertRaises(asyncio.CancelledError):
+                await self._get_server_session().send_request('blockchain.block.header', [999])
+        self.assertTrue(any(("Interface.[127.0.0.1:" in msg and "unexpected request. not a notification" in msg)
+                            for msg in logs.output))
+
+    async def test_we_disconnect_on_incoming_notification_spam(self):
+        """We don't expect the server to send us any requests of its own."""
+        interface = await self._start_iface_and_wait_for_sync()
+        # set lower resource limits on client
+        assert interface.session.cost_hard_limit > 0
+        interface.session.cost_hard_limit /= 30
+        interface.session.cost_soft_limit = interface.session.cost_hard_limit - 1
+        interface.session.bw_cost_per_byte *= 10
+        with self.assertLogs('electrum', level='INFO') as logs:
+            # server now sends a crazy number of notifications to the client
+            srv_sess = self._get_server_session()
+            headersub_res = (srv_sess._get_headersub_result(),)
+            spam_count = 200_000
+            for i in range(spam_count):
+                if srv_sess.got_disconnected.is_set():
+                    break
+                await srv_sess.send_notification('blockchain.headers.subscribe', headersub_res)
+            # both parties should close their end of the session (triggered by the client DC-ing)
+            await asyncio.sleep(0)
+            assert srv_sess.got_disconnected.is_set()
+            async with util.async_timeout(1):
+                await interface.got_disconnected.wait()
+        # the client should not have received most of the spam:
+        assert interface.session.recv_count < spam_count / 20
+        assert interface.session.cost > interface.session.cost_hard_limit
+        self.assertTrue(any(("Interface.[127.0.0.1:" in msg and "closing session over resource usage" in msg)
+                            for msg in logs.output))

tests/toyserver/toyserver.py

@@ -457,12 +457,14 @@ def __init__(self, *args, toyserver: ToyServer, **kwargs):
         self.subbed_scripthashes = set()  # type: set[str]
         self._method_counts = collections.defaultdict(int)  # type: dict[str, int]
         self.client_name = None
+        self.got_disconnected = asyncio.Event()
         self.svr.sessions.add(self)
 
     async def connection_lost(self):
         await super().connection_lost()
         self.logger.debug(f'{self.remote_address()} disconnected')
         self.svr.sessions.discard(self)
+        self.got_disconnected.set()
 
     async def handle_request(self, request):
         handlers = {