review comments

Author: arvidn

chia/_tests/core/server/test_rate_limits.py

@@ -68,14 +68,14 @@ async def test_limits_v2(incoming: bool, tx_msg: bool, limit_size: bool, monkeyp
     limits: dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]
 
     if limit_size:
-        agg_limit = RLSettings(False, count * 2, 1024, count * len(message_data))
+        agg_limit = RLSettings(False, count * 2, len(message_data), count * len(message_data))
     else:
-        agg_limit = RLSettings(False, count, 1024, count * 2 * len(message_data))
+        agg_limit = RLSettings(False, count, len(message_data), count * 2 * len(message_data))
 
     if limit_size:
-        limits = {msg_type: RLSettings(not tx_msg, count * 2, 1024, count * len(message_data))}
+        limits = {msg_type: RLSettings(not tx_msg, count * 2, len(message_data), count * len(message_data))}
     else:
-        limits = {msg_type: RLSettings(not tx_msg, count, 1024, count * 2 * len(message_data))}
+        limits = {msg_type: RLSettings(not tx_msg, count, len(message_data), count * 2 * len(message_data))}
 
     def mock_get_limits(
         our_capabilities: list[Capability], peer_capabilities: list[Capability]

chia/_tests/util/test_network_protocol_test.py

@@ -4,6 +4,8 @@
 import inspect
 from typing import Any, cast
 
+import pytest
+
 from chia.protocols import (
     farmer_protocol,
     full_node_protocol,
@@ -264,3 +266,19 @@ def test_missing_messages() -> None:
     assert types_in_module(shared_protocol) == shared_msgs, (
         f"message types were added or removed from shared_protocol. {STANDARD_ADVICE}"
     )
+
+
+@pytest.mark.parametrize("version", [1, 2])
+def test_rate_limits_complete(version: int) -> None:
+    from chia.protocols.protocol_message_types import ProtocolMessageTypes
+    from chia.server.rate_limit_numbers import rate_limits
+
+    if version == 1:
+        composed = rate_limits[1]
+    elif version == 2:
+        composed = {
+            **rate_limits[1],
+            **rate_limits[2],
+        }
+
+    assert set(composed.keys()) == set(ProtocolMessageTypes)

chia/server/rate_limit_numbers.py

@@ -1,7 +1,6 @@
 # All of these rate limits scale with the number of transactions so the aggregate amounts are higher
 from __future__ import annotations
 
-import copy
 import dataclasses
 from typing import Optional, Union
 
@@ -33,7 +32,13 @@ class Unlimited:
     max_size: int  # Max size of each request
 
 
-aggregate_limit = RLSettings(False, 1000, 100 * 1024 * 1024 * 1024, 100 * 1024 * 1024 * 1024)
+# for the aggregate limit, not all fields of RLSettings are used. Only "frequency" and "max_total_size"
+aggregate_limit = RLSettings(
+    aggregate_limit=False,
+    frequency=1000,
+    max_size=0,
+    max_total_size=100 * 1024 * 1024,
+)
 
 
 def get_rate_limits_to_use(
@@ -46,8 +51,10 @@ def get_rate_limits_to_use(
         # Use V2 rate limits
         if 2 in compose_rate_limits_cache:
             return compose_rate_limits_cache[2], aggregate_limit
-        composed = copy.copy(rate_limits[1])
-        composed.update(rate_limits[2])
+        composed = {
+            **rate_limits[1],
+            **rate_limits[2],
+        }
         compose_rate_limits_cache[2] = composed
         return composed, aggregate_limit
     else:
@@ -111,10 +118,14 @@ def get_rate_limits_to_use(
         ProtocolMessageTypes.request_puzzle_solution: RLSettings(True, 1000, 100),
         ProtocolMessageTypes.respond_puzzle_solution: RLSettings(True, 1000, 1024 * 1024),
         ProtocolMessageTypes.reject_puzzle_solution: RLSettings(True, 1000, 100),
+        ProtocolMessageTypes.none_response: RLSettings(False, 500, 100),
         ProtocolMessageTypes.new_peak_wallet: RLSettings(True, 200, 300),
         ProtocolMessageTypes.request_block_header: RLSettings(True, 500, 100),
         ProtocolMessageTypes.respond_block_header: RLSettings(True, 500, 500 * 1024),
         ProtocolMessageTypes.reject_header_request: RLSettings(True, 500, 100),
+        ProtocolMessageTypes.request_block_headers: RLSettings(False, 5000, 100),
+        ProtocolMessageTypes.reject_block_headers: RLSettings(False, 1000, 100),
+        ProtocolMessageTypes.respond_block_headers: RLSettings(False, 5000, 2 * 1024 * 1024),
         ProtocolMessageTypes.request_removals: RLSettings(True, 500, 50 * 1024, 10 * 1024 * 1024),
         ProtocolMessageTypes.respond_removals: RLSettings(True, 500, 1024 * 1024, 10 * 1024 * 1024),
         ProtocolMessageTypes.reject_removals_request: RLSettings(True, 500, 100),
@@ -160,6 +171,12 @@ def get_rate_limits_to_use(
         ProtocolMessageTypes.respond_ses_hashes: RLSettings(True, 2000, 1 * 1024 * 1024),
         ProtocolMessageTypes.request_children: RLSettings(True, 2000, 1024 * 1024),
         ProtocolMessageTypes.respond_children: RLSettings(True, 2000, 1 * 1024 * 1024),
+        ProtocolMessageTypes.error: RLSettings(False, 50000, 100),
+        ProtocolMessageTypes.request_fee_estimates: RLSettings(True, 10, 100),
+        ProtocolMessageTypes.respond_fee_estimates: RLSettings(True, 10, 100),
+        ProtocolMessageTypes.solve: RLSettings(False, 120, 1024),
+        ProtocolMessageTypes.solution_response: RLSettings(False, 120, 1024),
+        ProtocolMessageTypes.partial_proofs: RLSettings(False, 120, 3 * 1024),
     },
     2: {
         ProtocolMessageTypes.request_block_header: RLSettings(False, 500, 100),
@@ -173,18 +190,13 @@ def get_rate_limits_to_use(
         ProtocolMessageTypes.reject_additions_request: RLSettings(False, 500, 100),
         ProtocolMessageTypes.reject_header_blocks: RLSettings(False, 1000, 100),
         ProtocolMessageTypes.respond_header_blocks: RLSettings(False, 5000, 2 * 1024 * 1024),
-        ProtocolMessageTypes.request_block_headers: RLSettings(False, 5000, 100),
-        ProtocolMessageTypes.reject_block_headers: RLSettings(False, 1000, 100),
-        ProtocolMessageTypes.respond_block_headers: RLSettings(False, 5000, 2 * 1024 * 1024),
         ProtocolMessageTypes.request_ses_hashes: RLSettings(False, 2000, 1 * 1024 * 1024),
         ProtocolMessageTypes.respond_ses_hashes: RLSettings(False, 2000, 1 * 1024 * 1024),
         ProtocolMessageTypes.request_children: RLSettings(False, 2000, 1024 * 1024),
         ProtocolMessageTypes.respond_children: RLSettings(False, 2000, 1 * 1024 * 1024),
         ProtocolMessageTypes.request_puzzle_solution: RLSettings(False, 5000, 100),
         ProtocolMessageTypes.respond_puzzle_solution: RLSettings(False, 5000, 1024 * 1024),
         ProtocolMessageTypes.reject_puzzle_solution: RLSettings(False, 5000, 100),
-        ProtocolMessageTypes.none_response: RLSettings(False, 500, 100),
-        ProtocolMessageTypes.error: RLSettings(False, 50000, 100),
         # These will have a lower cap since they don't scale with high TPS (NON_TX_FREQ)
         ProtocolMessageTypes.request_header_blocks: RLSettings(True, 5000, 100),
     },

chia/server/rate_limits.py

@@ -84,11 +84,10 @@ def process_msg_and_check(
         rate_limits, agg_limit = get_rate_limits_to_use(our_capabilities, peer_capabilities)
 
         try:
-            limits: Union[RLSettings, Unlimited]
-            limits = rate_limits[message_type]
+            limits: Union[RLSettings, Unlimited] = rate_limits[message_type]
             if isinstance(limits, RLSettings) and limits.aggregate_limit:
                 non_tx_freq = agg_limit.frequency
-                assert agg_limit.max_total_size
+                assert agg_limit.max_total_size is not None
                 non_tx_max_total_size = agg_limit.max_total_size
                 new_non_tx_count = self.non_tx_message_counts + 1
                 new_non_tx_size = self.non_tx_cumulative_size + len(message.data)