From dc04270c19ed12c48b4ff43706164f2e207871f4 Mon Sep 17 00:00:00 2001
From: unniznd <au6695@gmail.com>
Date: Fri, 15 Aug 2025 13:53:24 +0530
Subject: [PATCH] fix: message id type inonsistency in handle ihave and message
 id parsing improvement in handle iwant

---
 libp2p/custom_types.py     |  1 +
 libp2p/pubsub/gossipsub.py | 21 +++++++++++----------
 libp2p/pubsub/utils.py     | 31 +++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 10 deletions(-)
 create mode 100644 libp2p/pubsub/utils.py

diff --git a/libp2p/custom_types.py b/libp2p/custom_types.py
index 0b844133..00f86ee8 100644
--- a/libp2p/custom_types.py
+++ b/libp2p/custom_types.py
@@ -37,3 +37,4 @@ SyncValidatorFn = Callable[[ID, rpc_pb2.Message], bool]
 AsyncValidatorFn = Callable[[ID, rpc_pb2.Message], Awaitable[bool]]
 ValidatorFn = Union[SyncValidatorFn, AsyncValidatorFn]
 UnsubscribeFn = Callable[[], Awaitable[None]]
+MessageID = NewType("MessageID", str)
diff --git a/libp2p/pubsub/gossipsub.py b/libp2p/pubsub/gossipsub.py
index cebc438b..d396c776 100644
--- a/libp2p/pubsub/gossipsub.py
+++ b/libp2p/pubsub/gossipsub.py
@@ -1,6 +1,3 @@
-from ast import (
-    literal_eval,
-)
 from collections import (
     defaultdict,
 )
@@ -22,6 +19,7 @@ from libp2p.abc import (
     IPubsubRouter,
 )
 from libp2p.custom_types import (
+    MessageID,
     TProtocol,
 )
 from libp2p.peer.id import (
@@ -54,6 +52,10 @@ from .pb import (
 from .pubsub import (
     Pubsub,
 )
+from .utils import (
+    parse_message_id_safe,
+    safe_parse_message_id,
+)
 
 PROTOCOL_ID = TProtocol("/meshsub/1.0.0")
 PROTOCOL_ID_V11 = TProtocol("/meshsub/1.1.0")
@@ -780,11 +782,10 @@ class GossipSub(IPubsubRouter, Service):
 
         # Add all unknown message ids (ids that appear in ihave_msg but not in
         # seen_seqnos) to list of messages we want to request
-        # FIXME: Update type of message ID
-        msg_ids_wanted: list[Any] = [
-            msg_id
+        msg_ids_wanted: list[MessageID] = [
+            parse_message_id_safe(msg_id)
             for msg_id in ihave_msg.messageIDs
-            if literal_eval(msg_id) not in seen_seqnos_and_peers
+            if msg_id not in str(seen_seqnos_and_peers)
         ]
 
         # Request messages with IWANT message
@@ -798,9 +799,9 @@ class GossipSub(IPubsubRouter, Service):
         Forwards all request messages that are present in mcache to the
         requesting peer.
         """
-        # FIXME: Update type of message ID
-        # FIXME: Find a better way to parse the msg ids
-        msg_ids: list[Any] = [literal_eval(msg) for msg in iwant_msg.messageIDs]
+        msg_ids: list[tuple[bytes, bytes]] = [
+            safe_parse_message_id(msg) for msg in iwant_msg.messageIDs
+        ]
         msgs_to_forward: list[rpc_pb2.Message] = []
         for msg_id_iwant in msg_ids:
             # Check if the wanted message ID is present in mcache
diff --git a/libp2p/pubsub/utils.py b/libp2p/pubsub/utils.py
new file mode 100644
index 00000000..13961873
--- /dev/null
+++ b/libp2p/pubsub/utils.py
@@ -0,0 +1,31 @@
+import ast
+
+from libp2p.custom_types import (
+    MessageID,
+)
+
+
+def parse_message_id_safe(msg_id_str: str) -> MessageID:
+    """Safely handle message ID as string."""
+    return MessageID(msg_id_str)
+
+
+def safe_parse_message_id(msg_id_str: str) -> tuple[bytes, bytes]:
+    """
+    Safely parse message ID using ast.literal_eval with validation.
+    :param msg_id_str: String representation of message ID
+    :return: Tuple of (seqno, from_id) as bytes
+    :raises ValueError: If parsing fails
+    """
+    try:
+        parsed = ast.literal_eval(msg_id_str)
+        if not isinstance(parsed, tuple) or len(parsed) != 2:
+            raise ValueError("Invalid message ID format")
+
+        seqno, from_id = parsed
+        if not isinstance(seqno, bytes) or not isinstance(from_id, bytes):
+            raise ValueError("Message ID components must be bytes")
+
+        return (seqno, from_id)
+    except (ValueError, SyntaxError) as e:
+        raise ValueError(f"Invalid message ID format: {e}")