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

Code before:

# this class is used to indicate that a message type is not subject to a rate
# limit, but just a per-message size limit. This may be appropriate for response
# messages that are implicitly limited by their corresponding request message
# Unlimited message types are also not subject to the overall limit across all
# messages (just like messages in the "tx" category)
@dataclasses.dataclass(frozen=True)
class Unlimited:
    max_size: int  # Max size of each request


aggregate_limit = RLSettings(False, 1000, 100 * 1024 * 1024 * 1024, 100 * 1024 * 1024 * 1024)


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 = 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),

Review comment (suggestion):
i personally find literals to be tidy, clear, non-mutating, and very flexible.


```suggestion
        composed = {
            **rate_limits[1],
            **rate_limits[2],
        )
```

Code after:
# this class is used to indicate that a message type is not subject to a rate
# limit, but just a per-message size limit. This may be appropriate for response
# messages that are implicitly limited by their corresponding request message
# Unlimited message types are also not subject to the overall limit across all
# messages (just like messages in the "tx" category)
@dataclasses.dataclass(frozen=True)
class Unlimited:
    max_size: int  # Max size of each request


# 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(
    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