fix(networking): update message-id calculation (leanEthereum#17)

* fix(networking): update message-id calculation

* Update gossipsub.py

Co-authored-by: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>

* Update gossipsub.py

Co-authored-by: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>

* Make ruff happy

---------

Co-authored-by: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>

Author: syjn99

docs/client/networking.md

@@ -101,24 +101,20 @@ The optional `from` (1), `seqno` (3), `signature` (5) and `key` (6) protobuf
 fields are omitted from the message, since messages are identified by content,
 anonymous, and signed where necessary in the application layer.
 
-The `message-id` of a gossipsub message MUST be the following 20 byte value
-computed from the message data:
+The `message-id` MUST be the following 20 byte value computed from the message:
 
-- If `message.data` has a valid snappy decompression, set `message-id` to the
-  first 20 bytes of the `SHA256` hash of the concatenation of
-  `MESSAGE_DOMAIN_VALID_SNAPPY` with the snappy decompressed message data, i.e.
-  `SHA256(MESSAGE_DOMAIN_VALID_SNAPPY + snappy_decompress(message.data))[:20]`.
-- Otherwise, set `message-id` to the first 20 bytes of the `SHA256` hash of the
-  concatenation of `MESSAGE_DOMAIN_INVALID_SNAPPY` with the raw message data,
-  i.e. `SHA256(MESSAGE_DOMAIN_INVALID_SNAPPY + message.data)[:20]`.
+- If `message.data` has a valid snappy decompression, set `message-id` to the first 20 bytes of the `SHA256` hash of
+  the concatenation of the following data: `MESSAGE_DOMAIN_VALID_SNAPPY`, the length of the topic byte string (encoded as little-endian `uint64`),
+  the topic byte string, and the snappy decompressed message data:
+  i.e. `SHA256(MESSAGE_DOMAIN_VALID_SNAPPY + uint_to_bytes(uint64(len(message.topic))) + message.topic + snappy_decompress(message.data))[:20]`.
+- Otherwise, set `message-id` to the first 20 bytes of the `SHA256` hash of
+  the concatenation of the following data: `MESSAGE_DOMAIN_INVALID_SNAPPY`, the length of the topic byte string (encoded as little-endian `uint64`),
+  the topic byte string, and the raw message data:
+  i.e. `SHA256(MESSAGE_DOMAIN_INVALID_SNAPPY + uint_to_bytes(uint64(len(message.topic))) + message.topic + message.data)[:20]`.
 
 Where relevant, clients MUST reject messages with `message-id` sizes other than
 20 bytes.
 
-*Note*: The above logic handles two exceptional cases: (1) multiple snappy
-`data` can decompress to the same value, and (2) some message `data` can fail to
-snappy decompress altogether.
-
 The payload is carried in the `data` field of a gossipsub message, and varies
 depending on the topic:
 

src/lean_spec/subspecs/networking/gossipsub.py

@@ -1,6 +1,20 @@
-"""Gossipsub protocol parameters."""
+"""
+Gossipsub protocol
+
+- Parameters for gossipsub operation.
+- Message ID computation based on topic and message data, with snappy decompression handling.
+"""
+
+import hashlib
+from typing import Annotated, Callable, Optional
+
+from pydantic import Field
 
 from lean_spec.subspecs.chain.config import DEVNET_CONFIG
+from lean_spec.subspecs.networking.config import (
+    MESSAGE_DOMAIN_INVALID_SNAPPY,
+    MESSAGE_DOMAIN_VALID_SNAPPY,
+)
 from lean_spec.types import StrictBaseModel
 
 
@@ -46,3 +60,97 @@ class GossipsubParameters(StrictBaseModel):
 
     This is calculated as SECONDS_PER_SLOT * JUSTIFICATION_LOOKBACK_SLOTS * 2.
     """
+
+
+MessageId = Annotated[bytes, Field(min_length=20, max_length=20)]
+"""A 20-byte ID for gossipsub messages."""
+
+
+class GossipsubMessage:
+    """
+    Represents a gossipsub message and manages its ID computation.
+
+    This class encapsulates the topic, data, and the logic to generate a
+    message ID, correctly handling snappy decompression. The generated ID is
+    cached for efficiency.
+    """
+
+    def __init__(
+        self,
+        topic: bytes,
+        data: bytes,
+        snappy_decompress: Optional[Callable[[bytes], bytes]] = None,
+    ):
+        """
+        Initializes the message.
+
+        Args:
+            topic: The topic byte string.
+            data: The raw message data.
+            snappy_decompress: Optional snappy decompression function.
+        """
+        self.topic: bytes = topic
+        self.raw_data: bytes = data
+        self._snappy_decompress = snappy_decompress
+        # Cache for the computed ID
+        self._id: Optional[MessageId] = None
+
+    @property
+    def id(self) -> MessageId:
+        """
+        Computes and returns the 20-byte message ID.
+
+        The ID is computed on first access and then cached. The computation
+        logic depends on whether the message data can be successfully
+        decompressed with snappy.
+        """
+        # Return the cached ID if it's already been computed
+        if self._id is not None:
+            return self._id
+
+        domain: bytes
+        data_for_hash: bytes
+
+        if self._snappy_decompress:
+            try:
+                # Try to decompress the data with snappy
+                decompressed_data = self._snappy_decompress(self.raw_data)
+                # Valid snappy decompression - use valid domain
+                domain = MESSAGE_DOMAIN_VALID_SNAPPY
+                data_for_hash = decompressed_data
+            except Exception:
+                # Invalid snappy decompression - use invalid domain
+                domain = MESSAGE_DOMAIN_INVALID_SNAPPY
+                data_for_hash = self.raw_data
+        else:
+            # No decompressor provided - use invalid domain
+            domain = MESSAGE_DOMAIN_INVALID_SNAPPY
+            data_for_hash = self.raw_data
+
+        # The internal computation returns the raw bytes...
+        computed_id_bytes = self._compute_raw_id(domain, data_for_hash)
+
+        # We then cast to our strict NewType before caching and returning.
+        self._id = MessageId(computed_id_bytes)
+        return self._id
+
+    def _compute_raw_id(self, domain: bytes, message_data: bytes) -> bytes:
+        """
+        Computes SHA256(domain + uint64_le(len(topic)) + topic + message_data)[:20].
+
+        Args:
+            domain: The 4-byte domain for message-id isolation.
+            message_data: The message data (either decompressed or raw).
+
+        Returns:
+            A 20-byte raw bytes digest.
+        """
+        # Encode the topic length as little-endian bytes
+        topic_len_bytes = len(self.topic).to_bytes(8, "little")
+
+        # Concatenate all components for hashing
+        data_to_hash = domain + topic_len_bytes + self.topic + message_data
+
+        # Compute SHA256 and take the first 20 bytes
+        digest = hashlib.sha256(data_to_hash).digest()
+        return digest[:20]