From 2d6741c9c861b20114d5012cb1dc4c18fa1b3b0f Mon Sep 17 00:00:00 2001 From: arvidn Date: Tue, 30 Jun 2026 15:01:46 +0200 Subject: [PATCH] make parsing of ProofOfSpace more strict and less maleable, like the original Option<> parsing --- crates/chia-protocol/src/proof_of_space.rs | 121 ++++++++++++++++----- 1 file changed, 94 insertions(+), 27 deletions(-) diff --git a/crates/chia-protocol/src/proof_of_space.rs b/crates/chia-protocol/src/proof_of_space.rs index c4978261a..73da03104 100644 --- a/crates/chia-protocol/src/proof_of_space.rs +++ b/crates/chia-protocol/src/proof_of_space.rs @@ -279,7 +279,7 @@ impl Streamable for ProofOfSpace { let pool_public_key = as Streamable>::parse::(input)?; let prefix = ::parse::(input)?; - let version = u8::from((prefix & 0b10) != 0); + let version = prefix >> 1; let pool_contract_puzzle_hash = if (prefix & 1) != 0 { Some(::parse::(input)?) } else { @@ -353,6 +353,34 @@ mod tests { G1Element::from_bytes(&POOL_PK_BYTES).expect("POOL_PK_BYTES is valid") } + fn make_pos(version: u8, has_pool_pk: bool, has_contract: bool, size: u8) -> ProofOfSpace { + ProofOfSpace::new( + Bytes32::default(), + has_pool_pk.then(pool_pk), + has_contract.then(Bytes32::default), + plot_pk(), + version, + 0, + 0, + 0, + size, + Bytes::from(vec![0x80]), + ) + } + + // Locate the pool_contract_puzzle_hash Optional prefix byte (which also + // encodes the version) by finding where a contract-present and a + // contract-absent serialization first differ. + fn prefix_offset() -> usize { + let absent = make_pos(0, true, false, 32).to_bytes().unwrap(); + let present = make_pos(0, true, true, 32).to_bytes().unwrap(); + absent + .iter() + .zip(present.iter()) + .position(|(a, b)| a != b) + .unwrap() + } + // these are regression tests and test vectors for plot ID computations #[rstest] #[case("pool_pk", hex!("e185d4ec721ec060eb5833ec07d802fc69a43ed45dd59d7f20c58494421e0270"))] @@ -466,37 +494,76 @@ mod tests { assert_eq!(quality, Bytes32::new(expect_quality)); } + // Round-trip every accepted prefix-byte combination. For v0 (legacy v1 + // proofs) the pool_contract Optional is lenient and accepts any of + // {neither, pool_pk, contract, both}, exactly like the original plain + // Option<> derive. For v1 (v2 proofs) the size field is not stored and + // collapses to 0 on parse. #[rstest] - #[case(0, 18, Ok(18))] - #[case(0, 28, Ok(28))] - #[case(0, 38, Ok(38))] - #[case(1, 18, Ok(0))] - #[case(1, 28, Ok(0))] - #[case(1, 38, Ok(0))] - #[case(2, 18, Err(Error::InvalidPoS))] - fn proof_of_space_size(#[case] version: u8, #[case] size: u8, #[case] expect: Result) { - let pos = ProofOfSpace::new( - Bytes32::from(b"abababababababababababababababab"), - Some(G1Element::default()), - None, - G1Element::default(), - version, - 0, - 0, - 0, - size, - Bytes::from(vec![]), - ); + #[case::v0_pool(0, true, false, 18, 18)] + #[case::v0_contract(0, false, true, 28, 28)] + #[case::v0_both_lenient(0, true, true, 38, 38)] + #[case::v0_neither_lenient(0, false, false, 10, 10)] + #[case::v1_pool(1, true, false, 18, 0)] + #[case::v1_contract(1, false, true, 28, 0)] + fn roundtrip( + #[case] version: u8, + #[case] has_pool_pk: bool, + #[case] has_contract: bool, + #[case] size: u8, + #[case] expected_size: u8, + ) { + let buf = make_pos(version, has_pool_pk, has_contract, size) + .to_bytes() + .unwrap(); + let parsed = ProofOfSpace::parse::(&mut Cursor::<&[u8]>::new(&buf)).unwrap(); + + assert_eq!(parsed.version, version); + assert_eq!(parsed.pool_public_key.is_some(), has_pool_pk); + assert_eq!(parsed.pool_contract_puzzle_hash.is_some(), has_contract); + assert_eq!(parsed.size, expected_size); + // re-serialization must be byte-for-byte stable + assert_eq!(parsed.to_bytes().unwrap(), buf); + } + // Inputs that must be rejected. The v2 (version 1) format requires exactly + // one of pool_public_key / pool_contract_puzzle_hash to be set, so both-set + // and neither-set are rejected at parse time. An unknown version (>= 2) is + // rejected at serialize time. The exact error code is unimportant, only + // that it fails. + #[rstest] + #[case::v1_both_set(1, true, true)] + #[case::v1_neither_set(1, false, false)] + #[case::unknown_version(2, true, false)] + fn parse_rejects(#[case] version: u8, #[case] has_pool_pk: bool, #[case] has_contract: bool) { + let pos = make_pos(version, has_pool_pk, has_contract, 0); match pos.to_bytes() { Ok(buf) => { - let new_pos = - ProofOfSpace::parse::(&mut Cursor::<&[u8]>::new(&buf)).expect("parse()"); - assert_eq!(new_pos.size, expect.unwrap()); - } - Err(e) => { - assert_eq!(e, expect.unwrap_err()); + let err = ProofOfSpace::parse::(&mut Cursor::<&[u8]>::new(&buf)) + .expect_err("must fail to parse"); + assert_eq!(err, Error::InvalidPoS); } + Err(e) => assert_eq!(e, Error::InvalidPoS), } } + + // The version flag is packed into the pool_contract_puzzle_hash Optional + // prefix byte. Only bit 0 (the Optional flag) and bit 1 (the version) carry + // meaning. The high bits (2..=7) must be rejected for both versions, + // matching the strictness of a plain Optional<> prefix in earlier protocol + // versions where this byte could only ever be 0 or 1. + #[rstest] + fn high_prefix_bits_rejected(#[values(0, 1)] version: u8, #[values(2, 3, 4, 5, 6, 7)] bit: u8) { + let offset = prefix_offset(); + let size = if version == 0 { 32 } else { 0 }; + let mut buf = make_pos(version, true, false, size).to_bytes().unwrap(); + + let expected_prefix = if version == 0 { 0b00 } else { 0b10 }; + assert_eq!(buf[offset], expected_prefix); + + buf[offset] |= 1u8 << bit; + let err = ProofOfSpace::parse::(&mut Cursor::<&[u8]>::new(&buf)) + .expect_err("high prefix bit must be rejected"); + assert_eq!(err, Error::InvalidPoS); + } }