[CHIA-4259] Short-circuit on invalid VDF data (#21196)

chia/_tests/core/custom_types/test_vdf.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import pytest
+from chia_rs import VDFInfo, VDFProof
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint8, uint64
+
+from chia.consensus.default_constants import DEFAULT_CONSTANTS
+from chia.types.blockchain_format import vdf
+from chia.types.blockchain_format.classgroup import ClassgroupElement
+
+
+@pytest.mark.parametrize(
+    ("witness_type", "witness_size"),
+    [
+        (uint8(0), 100),
+        (uint8(1), 241),
+        (uint8(63), 8983),
+    ],
+)
+def test_expected_witness_sizes_reach_verifier(
+    monkeypatch: pytest.MonkeyPatch, witness_type: uint8, witness_size: int
+) -> None:
+    verifier_calls = 0
+
+    def accept_proof(*args: object) -> bool:
+        nonlocal verifier_calls
+        verifier_calls += 1
+        return True
+
+    monkeypatch.setattr(vdf, "get_discriminant", lambda *args: 1)
+    monkeypatch.setattr(vdf, "verify_vdf", accept_proof)
+
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    proof = VDFProof(witness_type, bytes(witness_size), False)
+
+    assert vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)
+    assert verifier_calls == 1
+
+
+def test_oversized_witness_rejected_before_cached_verifier() -> None:
+    cache_info_before = vdf.verify_vdf.cache_info()
+
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    proof = VDFProof(uint8(0), bytes(2_000_000), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)
+    assert vdf.verify_vdf.cache_info() == cache_info_before
+
+
+@pytest.mark.parametrize(
+    ("witness_type", "witness_size"),
+    [
+        # Empty / truncated / one-byte-long for witness_type 0 (expected 100).
+        (uint8(0), 0),
+        (uint8(0), 99),
+        (uint8(0), 101),
+        # Correct length for a different witness_type.
+        (uint8(1), 100),
+        (uint8(0), 241),
+        # Off-by-one around witness_type 1 (expected 241).
+        (uint8(1), 240),
+        (uint8(1), 242),
+        # Off-by-one around witness_type 63 (expected 8983).
+        (uint8(63), 8982),
+        (uint8(63), 8984),
+    ],
+    ids=[
+        "type0_empty",
+        "type0_one_short",
+        "type0_one_long",
+        "type1_with_type0_size",
+        "type0_with_type1_size",
+        "type1_one_short",
+        "type1_one_long",
+        "type63_one_short",
+        "type63_one_long",
+    ],
+)
+def test_invalid_witness_size_rejected(witness_type: uint8, witness_size: int) -> None:
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    proof = VDFProof(witness_type, bytes(witness_size), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)
+
+
+def test_witness_type_above_max_rejected() -> None:
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    # MAX_VDF_WITNESS_SIZE is 64; witness_type + 1 must be <= 64, so 64 is rejected.
+    proof = VDFProof(uint8(64), bytes(100), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)
+
+
+def test_target_vdf_info_mismatch_rejected() -> None:
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    target = VDFInfo(bytes32(b"\x01" * 32), uint64(1), classgroup_element)
+    proof = VDFProof(uint8(0), bytes(100), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info, target)
+
+
+def test_verifier_failure_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(vdf, "get_discriminant", lambda *args: 1)
+    monkeypatch.setattr(vdf, "verify_vdf", lambda *args: False)
+
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    proof = VDFProof(uint8(0), bytes(100), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)
+
+
+def test_verifier_exception_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(vdf, "get_discriminant", lambda *args: 1)
+
+    def raise_error(*args: object) -> bool:
+        raise RuntimeError("chiavdf blew up")
+
+    monkeypatch.setattr(vdf, "verify_vdf", raise_error)
+
+    classgroup_element = ClassgroupElement.get_default_element()
+    info = VDFInfo(bytes32.zeros, uint64(1), classgroup_element)
+    proof = VDFProof(uint8(0), bytes(100), False)
+
+    assert not vdf.validate_vdf(proof, DEFAULT_CONSTANTS, classgroup_element, info)

chia/types/blockchain_format/vdf.py

@@ -61,6 +61,14 @@ def validate_vdf(
         return False
     if proof.witness_type + 1 > constants.MAX_VDF_WITNESS_SIZE:
         return False
+    # A witness holds one serialized classgroup element followed by one segment per
+    # level of recursion. Each segment stores an iteration count, a 264-bit challenge
+    # prime (33 bytes), and another classgroup element.
+    form_size = ClassgroupElement.get_size()
+    witness_segment_size = uint64.SIZE + 33 + form_size
+    expected_witness_size = form_size + proof.witness_type * witness_segment_size
+    if len(proof.witness) != expected_witness_size:
+        return False
     if len(input_el.data) != 100:
         log.error(f"Invalid ClassgroupElement size: {len(input_el.data)} (expected 100)")
         return False