diff --git a/minichain/chain.py b/minichain/chain.py index fa7d0dd..0eeae48 100644 --- a/minichain/chain.py +++ b/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, [] diff --git a/tests/test_pow_validation.py b/tests/test_pow_validation.py new file mode 100644 index 0000000..77695a5 --- /dev/null +++ b/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