fix: validate proof of work on block acceptance

minichain/chain.py

@@ -10,6 +10,15 @@
 logger = logging.getLogger(__name__)
 
 
+class InvalidProofOfWorkError(ValueError):
+    pass
+
+
+def validate_difficulty(difficulty, max_difficulty):
+    if type(difficulty) is not int or difficulty < 1 or difficulty > max_difficulty:
+        raise ValueError(f"invalid difficulty {difficulty}")
+
+
 def validate_block_link_and_hash(previous_block, block):
     if block.previous_hash != previous_block.hash:
         raise ValueError(
@@ -25,6 +34,12 @@ def validate_block_link_and_hash(previous_block, block):
     if block.hash != expected_hash:
         raise ValueError(f"invalid hash {block.hash}")
 
+    validate_difficulty(block.difficulty, len(block.hash))
+
+    if not block.hash.startswith("0" * block.difficulty):
+        raise InvalidProofOfWorkError(
+            f"invalid PoW: hash {block.hash} does not satisfy difficulty {block.difficulty}"
+        )
 
 class Blockchain:
     """
@@ -118,7 +133,11 @@ def get_total_work(self, chain_list=None):
         if chain_list is None:
             with self._lock:
                 chain_list = self.chain
-        return sum(2 ** (block.difficulty or 1) for block in chain_list)
+
+        for block in chain_list:
+            validate_difficulty(block.difficulty, len(block.hash))
+
+        return sum(2 ** block.difficulty for block in chain_list)
 
     def _next_difficulty(self, difficulty, avg_block_time):
         """Advance the EMA difficulty control after a block, returning the new value."""
@@ -139,11 +158,13 @@ def _apply_block(self, prev_block, block, state, difficulty, avg_block_time):
 
         try:
             validate_block_link_and_hash(prev_block, block)
+        except InvalidProofOfWorkError as exc:
+            logger.warning("Block %s rejected: %s", block.index, exc)
+            return ValidationStatus.INVALID, difficulty, avg_block_time
         except ValueError as exc:
             logger.warning("Block %s rejected: %s", block.index, exc)
             status = ValidationStatus.INVALID if "hash" in str(exc) else ValidationStatus.FAILED
             return status, difficulty, avg_block_time
-
         if block.difficulty != difficulty:
             logger.warning("Block %s rejected: Invalid difficulty. Expected %s, got %s", block.index, difficulty, block.difficulty)
             return ValidationStatus.INVALID, difficulty, avg_block_time
@@ -212,9 +233,12 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
             return False, []
 
         with self._lock:
-            current_work = self.get_total_work()
-            new_work = self.get_total_work(new_chain_list)
-
+            try:
+                current_work = self.get_total_work()
+                new_work = self.get_total_work(new_chain_list)
+            except ValueError as exc:
+                logger.warning("Reorg rejected: %s", exc)
+                return False, []
             if new_work <= current_work:
                 logger.debug("Incoming chain (work: %s) is not heavier than local chain (work: %s). Rejecting.", new_work, current_work)
                 return False, []

tests/test_pow_validation.py

@@ -0,0 +1,36 @@
+from minichain import Blockchain, Block
+from minichain.validators import ValidationStatus
+
+
+def _make_block_with_invalid_pow(chain):
+    block = Block(
+        index=chain.last_block.index + 1,
+        previous_hash=chain.last_block.hash,
+        transactions=[],
+        difficulty=chain.current_difficulty,
+        state_root=chain.state.state_root(),
+    )
+
+    while True:
+        block.hash = block.compute_hash()
+        if not block.hash.startswith("0" * block.difficulty):
+            return block
+        block.nonce += 1
+
+
+def test_add_block_rejects_invalid_pow():
+    chain = Blockchain()
+    block = _make_block_with_invalid_pow(chain)
+
+    assert chain.add_block(block) == ValidationStatus.INVALID
+    assert len(chain.chain) == 1
+
+
+def test_resolve_conflicts_rejects_invalid_pow():
+    chain = Blockchain()
+    block = _make_block_with_invalid_pow(chain)
+
+    success, _ = chain.resolve_conflicts([chain.chain[0], block])
+
+    assert success is False
+    assert len(chain.chain) == 1