feat: implement Merkle Patricia Trie for state verification

main.py

@@ -76,18 +76,21 @@ def mine_and_process_block(chain, mempool, miner_pk):
         logger.info("No mineable transactions in current queue window.")
         return None
 
+    temp_state.credit_mining_reward(miner_pk)
+
     block = Block(
         index=chain.last_block.index + 1,
         previous_hash=chain.last_block.hash,
         transactions=mineable_txs,
+        state_root=temp_state.state_root(),
+        miner=miner_pk,
     )
 
     mined_block = mine_block(block)
 
     if chain.add_block(mined_block):
         logger.info("✅ Block #%d mined and added (%d txs)", mined_block.index, len(mineable_txs))
         mempool.remove_transactions(mineable_txs)
-        chain.state.credit_mining_reward(miner_pk)
         return mined_block
     else:
         logger.error("❌ Block rejected by chain")
@@ -146,10 +149,6 @@ async def handler(data):
             if chain.add_block(block):
                 logger.info("📥 Received Block #%d — added to chain", block.index)
 
-                # Apply mining reward for the remote miner (burn address as placeholder)
-                miner = payload.get("miner", BURN_ADDRESS)
-                chain.state.credit_mining_reward(miner)
-
                 # Drop only confirmed transactions so higher nonces can remain queued.
                 mempool.remove_transactions(block.transactions)
             else:

minichain/block.py

@@ -41,6 +41,8 @@ def __init__(
         transactions: Optional[List[Transaction]] = None,
         timestamp: Optional[float] = None,
         difficulty: Optional[int] = None,
+        state_root: Optional[str] = None,
+        miner: Optional[str] = None,
     ):
         self.index = index
         self.previous_hash = previous_hash
@@ -56,6 +58,8 @@ def __init__(
         self.difficulty: Optional[int] = difficulty
         self.nonce: int = 0
         self.hash: Optional[str] = None
+        self.state_root: Optional[str] = state_root
+        self.miner: Optional[str] = miner
 
         # NEW: compute merkle root once
         self.merkle_root: Optional[str] = _calculate_merkle_root(self.transactions)
@@ -68,9 +72,11 @@ def to_header_dict(self):
             "index": self.index,
             "previous_hash": self.previous_hash,
             "merkle_root": self.merkle_root,
+            "state_root": self.state_root,
             "timestamp": self.timestamp,
             "difficulty": self.difficulty,
             "nonce": self.nonce,
+            "miner": self.miner,
         }
 
     # -------------------------
@@ -111,6 +117,8 @@ def from_dict(cls, payload: dict):
             transactions=transactions,
             timestamp=payload.get("timestamp"),
             difficulty=payload.get("difficulty"),
+            state_root=payload.get("state_root"),
+            miner=payload.get("miner"),
         )
         block.nonce = payload.get("nonce", 0)
         block.hash = payload.get("hash")

minichain/chain.py

@@ -71,7 +71,8 @@ def _create_genesis_block(self, genesis_path):
             previous_hash="0",
             transactions=[],
             timestamp=timestamp,
-            difficulty=difficulty
+            difficulty=difficulty,
+            state_root=self.state.state_root()
         )
 
         computed_hash = calculate_hash(genesis_block.to_header_dict())
@@ -119,6 +120,14 @@ def add_block(self, block):
                     logger.warning("Block %s rejected: Transaction failed validation", block.index)
                     return False
 
+            if block.miner:
+                temp_state.credit_mining_reward(block.miner)
+
+            # Verify state root
+            if block.state_root != temp_state.state_root():
+                logger.warning("Block %s rejected: Invalid state root. Expected %s, got %s", block.index, temp_state.state_root(), block.state_root)
+                return False
+
             # All transactions valid → commit state and append block
             self.state = temp_state
             self.chain.append(block)

minichain/mpt.py

@@ -0,0 +1,164 @@
+import hashlib
+import json
+from typing import Optional, List
+
+def hash_data(data: bytes) -> bytes:
+    return hashlib.sha256(data).digest()
+
+def to_nibbles(key_hex: str) -> List[int]:
+    """Converts a hex string key into a list of integer nibbles (0-15)."""
+    try:
+        return [int(c, 16) for c in key_hex]
+    except ValueError:
+        raise ValueError(f"Invalid MPT key: '{key_hex}'. Keys must be valid hex strings.")
+
+class Node:
+    def hash(self) -> bytes:
+        raise NotImplementedError
+
+class LeafNode(Node):
+    def __init__(self, path: List[int], value: str):
+        self.path = path
+        self.value = value
+
+    def hash(self) -> bytes:
+        data = json.dumps({"type": "leaf", "path": self.path, "value": self.value}, sort_keys=True)
+        return hash_data(data.encode())
+
+class ExtensionNode(Node):
+    def __init__(self, path: List[int], child: Node):
+        self.path = path
+        self.child = child
+
+    def hash(self) -> bytes:
+        child_hash = self.child.hash().hex()
+        data = json.dumps({"type": "extension", "path": self.path, "child": child_hash}, sort_keys=True)
+        return hash_data(data.encode())
+
+class BranchNode(Node):
+    def __init__(self):
+        self.branches: List[Optional[Node]] = [None] * 16
+        self.value: Optional[str] = None
+
+    def hash(self) -> bytes:
+        b_hashes = [b.hash().hex() if b else None for b in self.branches]
+        data = json.dumps({"type": "branch", "branches": b_hashes, "value": self.value}, sort_keys=True)
+        return hash_data(data.encode())
+
+class Trie:
+    """
+    A simplified Merkle Patricia Trie (MPT) for MiniChain.
+    Provides O(log N) state verification via cryptographic state roots.
+    """
+    def __init__(self):
+        self.root: Optional[Node] = None
+
+    def root_hash(self) -> str:
+        """Returns the 32-byte hex hash of the trie's root."""
+        if not self.root:
+            return "0" * 64
+        return self.root.hash().hex()
+
+    def get(self, key_hex: str) -> Optional[str]:
+        if not self.root:
+            return None
+        return self._get(self.root, to_nibbles(key_hex))
+
+    def _get(self, node: Optional[Node], path: List[int]) -> Optional[str]:
+        if not node:
+            return None
+
+        if isinstance(node, LeafNode):
+            if node.path == path:
+                return node.value
+            return None
+
+        elif isinstance(node, ExtensionNode):
+            if path[:len(node.path)] == node.path:
+                return self._get(node.child, path[len(node.path):])
+            return None
+
+        elif isinstance(node, BranchNode):
+            if not path:
+                return node.value
+            nibble = path[0]
+            return self._get(node.branches[nibble], path[1:])
+
+        return None
+
+    def put(self, key_hex: str, value: str):
+        path = to_nibbles(key_hex)
+        self.root = self._put(self.root, path, value)
+
+    def _put(self, node: Optional[Node], path: List[int], value: str) -> Node:
+        if node is None:
+            return LeafNode(path, value)
+
+        if isinstance(node, LeafNode):
+            if node.path == path:
+                node.value = value
+                return node
+
+            # Paths diverge. Find common prefix.
+            common = 0
+            while common < len(node.path) and common < len(path) and node.path[common] == path[common]:
+                common += 1
+
+            branch = BranchNode()
+
+            # Handle the leaf's remaining path
+            leaf_remaining = node.path[common:]
+            if not leaf_remaining:
+                branch.value = node.value
+            else:
+                branch.branches[leaf_remaining[0]] = LeafNode(leaf_remaining[1:], node.value)
+
+            # Handle the new value's remaining path
+            new_remaining = path[common:]
+            if not new_remaining:
+                branch.value = value
+            else:
+                branch.branches[new_remaining[0]] = LeafNode(new_remaining[1:], value)
+
+            if common > 0:
+                return ExtensionNode(node.path[:common], branch)
+            return branch
+
+        elif isinstance(node, ExtensionNode):
+            common = 0
+            while common < len(node.path) and common < len(path) and node.path[common] == path[common]:
+                common += 1
+
+            if common == len(node.path):
+                # Path matches extension exactly, continue to child
+                node.child = self._put(node.child, path[common:], value)
+                return node
+
+            # Divergence inside the extension node
+            branch = BranchNode()
+            ext_remaining = node.path[common:]
+
+            # The child of the extension becomes a branch's branch
+            if len(ext_remaining) == 1:
+                branch.branches[ext_remaining[0]] = node.child
+            else:
+                branch.branches[ext_remaining[0]] = ExtensionNode(ext_remaining[1:], node.child)
+
+            # Insert the new value
+            new_remaining = path[common:]
+            if not new_remaining:
+                branch.value = value
+            else:
+                branch.branches[new_remaining[0]] = LeafNode(new_remaining[1:], value)
+
+            if common > 0:
+                return ExtensionNode(node.path[:common], branch)
+            return branch
+
+        elif isinstance(node, BranchNode):
+            if not path:
+                node.value = value
+            else:
+                nibble = path[0]
+                node.branches[nibble] = self._put(node.branches[nibble], path[1:], value)
+            return node

minichain/p2p.py

@@ -183,23 +183,27 @@ def _validate_block_payload(self, payload):
             "index": int,
             "previous_hash": str,
             "merkle_root": (str, type(None)),
+            "state_root": str,
             "transactions": list,
             "timestamp": int,
             "difficulty": (int, type(None)),
             "nonce": int,
             "hash": str,
         }
-        optional_fields = {"miner": str}
+        optional_fields = {"miner": (str, type(None))}
         allowed_fields = set(required_fields) | set(optional_fields)
 
+        if not set(required_fields).issubset(payload):
+            return False
+
         if not set(payload).issubset(allowed_fields):
             return False
 
         for field, expected_type in required_fields.items():
             if not isinstance(payload.get(field), expected_type):
                 return False
 
-        if "miner" in payload and not isinstance(payload["miner"], str):
+        if "miner" in payload and not isinstance(payload["miner"], (str, type(None))):
             return False
 
         return all(

minichain/state.py

@@ -13,6 +13,21 @@ def __init__(self):
         self.accounts = {}
         self.contract_machine = ContractMachine(self)
 
+    def state_root(self) -> str:
+        """
+        Dynamically builds the Merkle Patricia Trie from the current state dictionary
+        and returns the cryptographic state root hash.
+        """
+        import json
+        from .mpt import Trie
+        trie = Trie()
+        # Sort items to ensure deterministic insertion order if necessary (though MPT is order-independent)
+        for addr, acc in sorted(self.accounts.items()):
+            if acc.get('balance', 0) == 0 and acc.get('nonce', 0) == 0 and not acc.get('code') and not acc.get('storage'):
+                continue
+            trie.put(addr, json.dumps(acc, sort_keys=True))
+        return trie.root_hash()
+
     DEFAULT_MINING_REWARD = 50
 
     def get_account(self, address):

tests/test_persistence.py

@@ -41,11 +41,15 @@ def _chain_with_tx(self):
         tx = Transaction(alice_pk, bob_pk, 30, 0)
         tx.sign(alice_sk)
 
+        temp_state = bc.state.copy()
+        temp_state.validate_and_apply(tx)
+
         block = Block(
             index=1,
             previous_hash=bc.last_block.hash,
             transactions=[tx],
             difficulty=1,
+            state_root=temp_state.state_root(),
         )
         mine_block(block, difficulty=1)
         bc.add_block(block)
@@ -216,11 +220,15 @@ def test_loaded_chain_can_add_new_block(self):
         tx2 = Transaction(new_pk, bob_pk, 10, 0)
         tx2.sign(new_sk)
 
+        temp_state = restored.state.copy()
+        temp_state.validate_and_apply(tx2)
+
         block2 = Block(
             index=len(restored.chain),
             previous_hash=restored.last_block.hash,
             transactions=[tx2],
             difficulty=1,
+            state_root=temp_state.state_root(),
         )
         mine_block(block2, difficulty=1)
 

tests/test_protocol_hardening.py

@@ -108,7 +108,7 @@ async def test_block_schema_accepts_current_block_wire_format(self):
         tx = Transaction(sender_pk, receiver_pk, 1, 0, timestamp=123)
         tx.sign(sender_sk)
 
-        block = Block(index=1, previous_hash="0" * 64, transactions=[tx], timestamp=456, difficulty=2)
+        block = Block(index=1, previous_hash="0" * 64, transactions=[tx], timestamp=456, difficulty=2, state_root="0"*64)
         block.nonce = 9
         block.hash = block.compute_hash()