Repository: Chia-Network/chia-blockchain
File: chia/server/rate_limit_numbers.py (Python)

Code before:
    # This will use the newest possible rate limits that both peers support. At this time there are only two
    # options, v1 and v2.

    if Capability.RATE_LIMITS_V2 in our_capabilities and Capability.RATE_LIMITS_V2 in peer_capabilities:
        # 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])
        compose_rate_limits_cache[2] = composed
        return composed, aggregate_limit
    else:
        # Use V1 rate limits
        return rate_limits[1], aggregate_limit


# Each number in this dict corresponds to a specific version of rate limits (1, 2,  etc).
# Version 1 includes the original limits for chia software from versions 1.0 to 1.4.
rate_limits: dict[int, dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]] = {
    1: {
        ProtocolMessageTypes.new_transaction: RLSettings(False, 5000, 100, 5000 * 100),
        ProtocolMessageTypes.request_transaction: RLSettings(False, 5000, 100, 5000 * 100),
        # TODO: check this
        ProtocolMessageTypes.respond_transaction: RLSettings(False, 5000, 1 * 1024 * 1024, 20 * 1024 * 1024),
        ProtocolMessageTypes.send_transaction: RLSettings(False, 5000, 1024 * 1024),
        ProtocolMessageTypes.transaction_ack: RLSettings(False, 5000, 2048),
        # All non-transaction apis also have an aggregate limit
        ProtocolMessageTypes.handshake: RLSettings(True, 5, 10 * 1024, 5 * 10 * 1024),
        ProtocolMessageTypes.harvester_handshake: RLSettings(True, 5, 1024 * 1024),
        ProtocolMessageTypes.new_signage_point_harvester: RLSettings(True, 100, 4886),  # Size with 100 pool list
        ProtocolMessageTypes.new_proof_of_space: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.request_signatures: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.respond_signatures: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.new_signage_point: RLSettings(True, 200, 2048),
        ProtocolMessageTypes.declare_proof_of_space: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.request_signed_values: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.farming_info: RLSettings(True, 100, 1024),
        ProtocolMessageTypes.signed_values: RLSettings(True, 100, 1024),
        ProtocolMessageTypes.new_peak_timelord: RLSettings(True, 100, 20 * 1024),
        ProtocolMessageTypes.new_unfinished_block_timelord: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.new_signage_point_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.new_infusion_point_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.new_end_of_sub_slot_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.request_compact_proof_of_time: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.respond_compact_proof_of_time: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.new_peak: RLSettings(True, 200, 512),
        ProtocolMessageTypes.request_proof_of_weight: RLSettings(True, 5, 100),
        ProtocolMessageTypes.respond_proof_of_weight: RLSettings(True, 5, 50 * 1024 * 1024, 100 * 1024 * 1024),
        ProtocolMessageTypes.request_block: RLSettings(True, 200, 100),
        ProtocolMessageTypes.reject_block: Unlimited(100),
        ProtocolMessageTypes.request_blocks: RLSettings(True, 500, 100),

Review comment (suggestion):
would be nice to have a test (in tests) maybe like...

```
@pytest.mark.parametrize(argnames="version", argvalues=[1, 2])
def test_rate_limits_complete(version: int) -> None:
    assert set(rate_limits[version].keys()) == set(ProtocolMessageTypes)
```

something about the 1 and 2 as well.  maybe also tieing them to the capability enumerators.  maybe.

Code after:


def get_rate_limits_to_use(
    our_capabilities: list[Capability], peer_capabilities: list[Capability]
) -> tuple[dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]], RLSettings]:
    # This will use the newest possible rate limits that both peers support. At this time there are only two
    # options, v1 and v2.

    if Capability.RATE_LIMITS_V2 in our_capabilities and Capability.RATE_LIMITS_V2 in peer_capabilities:
        # Use V2 rate limits
        if 2 in compose_rate_limits_cache:
            return compose_rate_limits_cache[2], aggregate_limit
        composed = {
            **rate_limits[1],
            **rate_limits[2],
        }
        compose_rate_limits_cache[2] = composed
        return composed, aggregate_limit
    else:
        # Use V1 rate limits
        return rate_limits[1], aggregate_limit


# Each number in this dict corresponds to a specific version of rate limits (1, 2,  etc).
# Version 1 includes the original limits for chia software from versions 1.0 to 1.4.
rate_limits: dict[int, dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]] = {
    1: {
        ProtocolMessageTypes.new_transaction: RLSettings(False, 5000, 100, 5000 * 100),
        ProtocolMessageTypes.request_transaction: RLSettings(False, 5000, 100, 5000 * 100),
        # TODO: check this
        ProtocolMessageTypes.respond_transaction: RLSettings(False, 5000, 1 * 1024 * 1024, 20 * 1024 * 1024),
        ProtocolMessageTypes.send_transaction: RLSettings(False, 5000, 1024 * 1024),
        ProtocolMessageTypes.transaction_ack: RLSettings(False, 5000, 2048),
        # All non-transaction apis also have an aggregate limit
        ProtocolMessageTypes.handshake: RLSettings(True, 5, 10 * 1024, 5 * 10 * 1024),
        ProtocolMessageTypes.harvester_handshake: RLSettings(True, 5, 1024 * 1024),
        ProtocolMessageTypes.new_signage_point_harvester: RLSettings(True, 100, 4886),  # Size with 100 pool list
        ProtocolMessageTypes.new_proof_of_space: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.request_signatures: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.respond_signatures: RLSettings(True, 100, 2048),
        ProtocolMessageTypes.new_signage_point: RLSettings(True, 200, 2048),
        ProtocolMessageTypes.declare_proof_of_space: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.request_signed_values: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.farming_info: RLSettings(True, 100, 1024),
        ProtocolMessageTypes.signed_values: RLSettings(True, 100, 1024),
        ProtocolMessageTypes.new_peak_timelord: RLSettings(True, 100, 20 * 1024),
        ProtocolMessageTypes.new_unfinished_block_timelord: RLSettings(True, 100, 10 * 1024),
        ProtocolMessageTypes.new_signage_point_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.new_infusion_point_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.new_end_of_sub_slot_vdf: RLSettings(True, 100, 100 * 1024),
        ProtocolMessageTypes.request_compact_proof_of_time: RLSettings(True, 100, 10 * 1024),