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

Code before:
        current_slot = int(self.get_time() // self.reset_seconds)
        if current_slot != self.current_slot:
            self.current_slot = current_slot
            self.message_counts = Counter()
            self.message_cumulative_sizes = Counter()
            self.non_tx_message_counts = 0
            self.non_tx_cumulative_size = 0
        try:
            message_type = ProtocolMessageTypes(message.type)
        except Exception as e:
            log.warning(f"Invalid message: {message.type}, {e}")
            return None

        new_message_counts: int = self.message_counts[message_type] + 1
        new_cumulative_size: int = self.message_cumulative_sizes[message_type] + len(message.data)
        new_non_tx_count: int = self.non_tx_message_counts
        new_non_tx_size: int = self.non_tx_cumulative_size
        proportion_of_limit: float = self.percentage_of_limit / 100

        ret: bool = False
        rate_limits: dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]
        rate_limits, agg_limit = get_rate_limits_to_use(our_capabilities, peer_capabilities)

        try:
            limits: Union[RLSettings, Unlimited]
            limits = rate_limits[message_type]
            if isinstance(limits, RLSettings) and limits.aggregate_limit:
                non_tx_freq = agg_limit.frequency
                assert agg_limit.max_total_size
                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)
                if new_non_tx_count > non_tx_freq * proportion_of_limit:
                    return " ".join(
                        [
                            f"non-tx count: {new_non_tx_count}",
                            f"> {non_tx_freq * proportion_of_limit}",
                            f"(scale factor: {proportion_of_limit})",
                        ]
                    )
                if new_non_tx_size > non_tx_max_total_size * proportion_of_limit:
                    return " ".join(
                        [
                            f"non-tx size: {new_non_tx_size}",
                            f"> {non_tx_max_total_size * proportion_of_limit}",
                            f"(scale factor: {proportion_of_limit})",
                        ]
                    )

            if isinstance(limits, Unlimited):
                # this message type is not rate limited. This is used for

Review comment (suggestion):
maybe worth keeping around an error message?  like `limits = rate_limits.get(message_type)` and `if limits is None:`.

anyways, could at least be:

```suggestion
            limits: Union[RLSettings, Unlimited] = rate_limits[message_type]
```

or to avoid accidentally expanding the type at some point

```suggestion
            limits = rate_limits[message_type]
```

Code after:
        current_slot = int(self.get_time() // self.reset_seconds)
        if current_slot != self.current_slot:
            self.current_slot = current_slot
            self.message_counts = Counter()
            self.message_cumulative_sizes = Counter()
            self.non_tx_message_counts = 0
            self.non_tx_cumulative_size = 0
        try:
            message_type = ProtocolMessageTypes(message.type)
        except Exception as e:
            log.warning(f"Invalid message: {message.type}, {e}")
            return None

        new_message_counts: int = self.message_counts[message_type] + 1
        new_cumulative_size: int = self.message_cumulative_sizes[message_type] + len(message.data)
        new_non_tx_count: int = self.non_tx_message_counts
        new_non_tx_size: int = self.non_tx_cumulative_size
        proportion_of_limit: float = self.percentage_of_limit / 100

        ret: bool = False
        rate_limits: dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]
        rate_limits, agg_limit = get_rate_limits_to_use(our_capabilities, peer_capabilities)

        try:
            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 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)
                if new_non_tx_count > non_tx_freq * proportion_of_limit:
                    return " ".join(
                        [
                            f"non-tx count: {new_non_tx_count}",
                            f"> {non_tx_freq * proportion_of_limit}",
                            f"(scale factor: {proportion_of_limit})",
                        ]
                    )
                if new_non_tx_size > non_tx_max_total_size * proportion_of_limit:
                    return " ".join(
                        [
                            f"non-tx size: {new_non_tx_size}",
                            f"> {non_tx_max_total_size * proportion_of_limit}",
                            f"(scale factor: {proportion_of_limit})",
                        ]
                    )

            if isinstance(limits, Unlimited):
                # this message type is not rate limited. This is used for
                # response messages and must be combined with banning peers