From 9717a95466ca386b6582c02319bae8ff0a16fa83 Mon Sep 17 00:00:00 2001 From: siddhant Date: Fri, 6 Mar 2026 19:10:56 +0530 Subject: [PATCH 1/3] address CodeRabbit networking comments --- main.py | 52 +++++++++++++++++++++++++++++++++--------------- minichain/p2p.py | 29 +++++++++++++++++++++++---- 2 files changed, 61 insertions(+), 20 deletions(-) diff --git a/main.py b/main.py index 12577dd..0900c40 100644 --- a/main.py +++ b/main.py @@ -31,6 +31,8 @@ logger = logging.getLogger(__name__) BURN_ADDRESS = "0" * 40 +TRUSTED_PEERS = set() +LOCALHOST_PEERS = {"127.0.0.1", "::1", "localhost", "0:0:0:0:0:0:0:1"} # ────────────────────────────────────────────── @@ -68,6 +70,11 @@ def mine_and_process_block(chain, mempool, miner_pk): return mined_block else: logger.error("❌ Block rejected by chain") + restored = 0 + for tx in pending_txs: + if mempool.add_transaction(tx): + restored += 1 + logger.info("Mempool: Restored %d/%d txs after rejection", restored, len(pending_txs)) return None @@ -81,15 +88,23 @@ def make_network_handler(chain, mempool): async def handler(data): msg_type = data.get("type") payload = data.get("data") + peer_addr = data.get("_peer_addr", "unknown") if msg_type == "sync": + peer_host = peer_addr.rsplit(":", 1)[0] if ":" in peer_addr else peer_addr + is_trusted = peer_addr in TRUSTED_PEERS or peer_host in TRUSTED_PEERS + is_localhost = peer_host in LOCALHOST_PEERS + if chain.state.accounts and not (is_trusted or is_localhost): + logger.warning("🔒 Rejected sync from untrusted peer %s", peer_addr) + return + # Merge remote state into local state (for accounts we don't have yet) - remote_accounts = payload.get("accounts", {}) + remote_accounts = payload.get("accounts", {}) if isinstance(payload, dict) else {} for addr, acc in remote_accounts.items(): if addr not in chain.state.accounts: chain.state.accounts[addr] = acc logger.info("🔄 Synced account %s... (balance=%d)", addr[:12], acc.get("balance", 0)) - logger.info("🔄 State sync complete — %d accounts", len(chain.state.accounts)) + logger.info("🔄 Accepted state sync from %s — %d accounts", peer_addr, len(chain.state.accounts)) elif msg_type == "tx": tx = Transaction(**payload) @@ -134,15 +149,15 @@ async def handler(data): ╔════════════════════════════════════════════════╗ ║ MiniChain Commands ║ ╠════════════════════════════════════════════════╣ -║ balance — show all balances ║ -║ send — send coins ║ -║ mine — mine a block ║ -║ peers — show connected peers ║ -║ connect — connect to a peer ║ -║ address — show your public key ║ -║ chain — show chain summary ║ -║ help — show this help ║ -║ quit — shut down ║ +║ balance - show all balances ║ +║ send - send coins ║ +║ mine - mine a block ║ +║ peers - show connected peers ║ +║ connect - connect to a peer ║ +║ address - show your public key ║ +║ chain - show chain summary ║ +║ help - show this help ║ +║ quit - shut down ║ ╚════════════════════════════════════════════════╝ """ @@ -220,7 +235,11 @@ async def cli_loop(sk, pk, chain, mempool, network, nonce_counter): except ValueError: print(" Invalid format. Use host:port") continue - await network.connect_to_peer(host, port) + success = await network.connect_to_peer(host, port) + if success: + print(f" Connected to {host}:{port}") + else: + print(f" Failed to connect to {host}:{port}") # ── address ── elif cmd == "address": @@ -249,7 +268,7 @@ async def cli_loop(sk, pk, chain, mempool, network, nonce_counter): # Main entry point # ────────────────────────────────────────────── -async def run_node(port: int, connect_to: str | None, fund: int): +async def run_node(port: int, connect_to: str | None, fund: int, host: str = "127.0.0.1"): """Boot the node, optionally connect to a peer, then enter the CLI.""" sk, pk = create_wallet() @@ -271,9 +290,9 @@ async def on_peer_connected(writer): await writer.drain() logger.info("🔄 Sent state sync to new peer") - network._on_peer_connected = on_peer_connected + network.set_on_peer_connected(on_peer_connected) - await network.start(port=port) + await network.start(port=port, host=host) # Fund this node's wallet so it can transact in the demo if fund > 0: @@ -299,6 +318,7 @@ async def on_peer_connected(writer): def main(): parser = argparse.ArgumentParser(description="MiniChain Node — Testnet Demo") + parser.add_argument("--host", type=str, default="127.0.0.1", help="Host/IP to bind the P2P server (default: 127.0.0.1)") parser.add_argument("--port", type=int, default=9000, help="TCP port to listen on (default: 9000)") parser.add_argument("--connect", type=str, default=None, help="Peer address to connect to (host:port)") parser.add_argument("--fund", type=int, default=100, help="Initial coins to fund this wallet (default: 100)") @@ -311,7 +331,7 @@ def main(): ) try: - asyncio.run(run_node(args.port, args.connect, args.fund)) + asyncio.run(run_node(args.port, args.connect, args.fund, args.host)) except KeyboardInterrupt: print("\nNode shut down.") diff --git a/minichain/p2p.py b/minichain/p2p.py index 81ff100..66737ec 100644 --- a/minichain/p2p.py +++ b/minichain/p2p.py @@ -37,23 +37,31 @@ def register_handler(self, handler_callback): raise ValueError("handler_callback must be callable") self._handler_callback = handler_callback + def set_on_peer_connected(self, callback): + if callback is not None and not asyncio.iscoroutinefunction(callback): + raise ValueError("on_peer_connected callback must be an async callable") + self._on_peer_connected = callback + # ------------------------------------------------------------------ # Server lifecycle # ------------------------------------------------------------------ - 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( - self._handle_incoming, "0.0.0.0", port + self._handle_incoming, host, port ) - logger.info("Network: Listening on 0.0.0.0:%d", port) + logger.info("Network: Listening on %s:%d", host, port) async def stop(self): """Gracefully shut down the server and disconnect all peers.""" logger.info("Network: Shutting down") for task in self._listen_tasks: task.cancel() + if self._listen_tasks: + await asyncio.gather(*self._listen_tasks, return_exceptions=True) + self._listen_tasks.clear() for _, writer in self._peers: try: writer.close() @@ -75,6 +83,11 @@ async def connect_to_peer(self, host: str, port: int) -> bool: try: reader, writer = await asyncio.open_connection(host, port) self._peers.append((reader, writer)) + if self._on_peer_connected: + try: + await self._on_peer_connected(writer) + except Exception: + logger.exception("Network: Error during peer sync") task = asyncio.create_task(self._listen_to_peer(reader, writer, f"{host}:{port}")) self._listen_tasks.append(task) logger.info("Network: Connected to peer %s:%d", host, port) @@ -110,6 +123,8 @@ async def _listen_to_peer(self, reader: asyncio.StreamReader, writer: asyncio.St except (json.JSONDecodeError, UnicodeDecodeError): logger.warning("Network: Malformed message from %s", addr) continue + if isinstance(data, dict): + data["_peer_addr"] = addr if self._handler_callback: try: @@ -144,7 +159,13 @@ async def _broadcast_raw(self, payload: dict): await writer.drain() except Exception: disconnected.append((reader, writer)) - for pair in disconnected: + for reader, writer in disconnected: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + pair = (reader, writer) if pair in self._peers: self._peers.remove(pair) From feb102651ed3213b1eef9154611859d4353144fe Mon Sep 17 00:00:00 2001 From: siddhant Date: Fri, 6 Mar 2026 19:27:27 +0530 Subject: [PATCH 2/3] fixed IPv6 peer --- main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/main.py b/main.py index 0900c40..1489e46 100644 --- a/main.py +++ b/main.py @@ -92,6 +92,7 @@ async def handler(data): if msg_type == "sync": peer_host = peer_addr.rsplit(":", 1)[0] if ":" in peer_addr else peer_addr + peer_host = peer_host.strip("[]") is_trusted = peer_addr in TRUSTED_PEERS or peer_host in TRUSTED_PEERS is_localhost = peer_host in LOCALHOST_PEERS if chain.state.accounts and not (is_trusted or is_localhost): From 6184794b028c606142b407f33924bba93d4528eb Mon Sep 17 00:00:00 2001 From: siddhant Date: Fri, 6 Mar 2026 19:55:22 +0530 Subject: [PATCH 3/3] harden sync payload checks and move initial funding after bootstrap --- main.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/main.py b/main.py index 1489e46..91fcc3c 100644 --- a/main.py +++ b/main.py @@ -100,8 +100,15 @@ async def handler(data): return # Merge remote state into local state (for accounts we don't have yet) - remote_accounts = payload.get("accounts", {}) if isinstance(payload, dict) else {} + remote_accounts = payload.get("accounts") if isinstance(payload, dict) else None + if not isinstance(remote_accounts, dict): + logger.warning("🔒 Rejected sync from %s with invalid accounts payload", peer_addr) + return + for addr, acc in remote_accounts.items(): + if not isinstance(acc, dict): + logger.warning("🔒 Skipping malformed account %r from %s", addr, peer_addr) + continue if addr not in chain.state.accounts: chain.state.accounts[addr] = acc logger.info("🔄 Synced account %s... (balance=%d)", addr[:12], acc.get("balance", 0)) @@ -154,7 +161,7 @@ async def handler(data): ║ send - send coins ║ ║ mine - mine a block ║ ║ peers - show connected peers ║ -║ connect - connect to a peer ║ +║ connect : - connect to a peer ║ ║ address - show your public key ║ ║ chain - show chain summary ║ ║ help - show this help ║ @@ -295,11 +302,6 @@ async def on_peer_connected(writer): await network.start(port=port, host=host) - # Fund this node's wallet so it can transact in the demo - if fund > 0: - chain.state.credit_mining_reward(pk, reward=fund) - logger.info("💰 Funded %s... with %d coins", pk[:12], fund) - # Connect to a seed peer if requested if connect_to: try: @@ -308,6 +310,11 @@ async def on_peer_connected(writer): except ValueError: logger.error("Invalid --connect format. Use host:port") + # Fund this node's wallet so it can transact in the demo + if fund > 0: + chain.state.credit_mining_reward(pk, reward=fund) + logger.info("💰 Funded %s... with %d coins", pk[:12], fund) + # Nonce counter kept as a mutable list so the CLI closure can mutate it nonce_counter = [0]