Skip to content

Fix P2P message validation and host parameter #64

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
quickquickcode wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix P2P message validation and host parameter #64

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ async def cli_loop(sk, pk, chain, mempool, network):
# Main entry point
# ──────────────────────────────────────────────

async def run_node(port: int, connect_to: str | None, fund: int, datadir: str | None):
async def run_node(port: int, host: str, connect_to: str | None, fund: int, datadir: str | None):
"""Boot the node, optionally connect to a peer, then enter the CLI."""
sk, pk = create_wallet()

Expand Down Expand Up @@ -326,7 +326,7 @@ async def on_peer_connected(writer):
await writer.drain()
logger.info("🔄 Sent state sync to new peer")

network.set_on_peer_connected(on_peer_connected)
network._on_peer_connected = on_peer_connected
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Prefer a public callback registration method over writing P2PNetwork internals.

Assigning network._on_peer_connected directly fixes the immediate error, but it hard-codes a private field name into main.py. A small public registration method in P2PNetwork would keep this contract centralized and easier to evolve.

Suggested API shape 
# minichain/p2p.py
+    def register_on_peer_connected(self, callback):
+        if not callable(callback):
+            raise ValueError("callback must be callable")
+        self._on_peer_connected = callback
# main.py
-    network._on_peer_connected = on_peer_connected
+    network.register_on_peer_connected(on_peer_connected)
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@main.py` at line 329, Replace the direct assignment to the private attribute
network._on_peer_connected by adding a public registration method on the
P2PNetwork class (e.g., register_on_peer_connected or
add_peer_connected_handler) that accepts the callback and sets/updates the
internal handler, then call that method from main.py instead of writing to
network._on_peer_connected; modify P2PNetwork to store the handler and invoke it
where peers connect, and update main.py to call the new public method with
on_peer_connected.


await network.start(port=port, host=host)

Expand Down Expand Up @@ -373,7 +373,7 @@ def main():
)

try:
asyncio.run(run_node(args.port, args.connect, args.fund, args.datadir))
asyncio.run(run_node(args.port, args.host, args.connect, args.fund, args.datadir))
except KeyboardInterrupt:
print("\nNode shut down.")

Expand Down
10 changes: 8 additions & 2 deletions minichain/p2p.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def register_handler(self, handler_callback):
raise ValueError("handler_callback must be callable")
self._handler_callback = handler_callback

async def start(self, port: int = 9000):
async def start(self, port: int = 9000, host: str = "127.0.0.1"):
"""Start listening for incoming peer connections on the given port."""
self._port = port
self._server = await asyncio.start_server(
Expand Down Expand Up @@ -206,7 +206,13 @@ def _validate_block_payload(self, payload):
def _validate_message(self, message):
if not isinstance(message, dict):
return False
if set(message) != {"type", "data"}:
# Allow _peer_addr field added by _listen_to_peer
required_fields = {"type", "data"}
if not required_fields.issubset(set(message)):
return False
# Reject messages with unexpected fields (except _peer_addr)
allowed_fields = {"type", "data", "_peer_addr"}
if not set(message).issubset(allowed_fields):
Comment on lines +209 to +215
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Keep _peer_addr out of the wire-schema validator.

_peer_addr is local transport metadata, not part of the documented on-the-wire message shape. A tighter fix is to validate the raw {"type", "data"} envelope first, then attach _peer_addr only after it passes.

Suggested fix 
     def _validate_message(self, message):
         if not isinstance(message, dict):
             return False
-        # Allow _peer_addr field added by _listen_to_peer
-        required_fields = {"type", "data"}
-        if not required_fields.issubset(set(message)):
-            return False
-        # Reject messages with unexpected fields (except _peer_addr)
-        allowed_fields = {"type", "data", "_peer_addr"}
-        if not set(message).issubset(allowed_fields):
+        if set(message) != {"type", "data"}:
             return False
-                if isinstance(data, dict):
-                    data["_peer_addr"] = addr
-
                 if not self._validate_message(data):
                     logger.warning("Network: Invalid message schema from %s", addr)
                     continue
+
+                data["_peer_addr"] = addr
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@minichain/p2p.py` around lines 209 - 215, The current validator is allowing
`_peer_addr` in the on‑wire schema; instead validate the raw envelope first by
checking required_fields = {"type","data"} against the incoming message (ensure
issubset of set(message)) and reject any extra fields immediately (i.e. ensure
set(message).issubset({"type","data"})) before attaching transport metadata;
remove `_peer_addr` from allowed_fields validation and only add `_peer_addr`
after the message passes validation (the place where `_listen_to_peer` attaches
`_peer_addr` should perform that attachment post‑validation).

return False

msg_type = message.get("type")
Expand Down
Loading