chia_protocol/
proof_of_space.rs

1use crate::bytes::{Bytes, Bytes32};
2use chia_bls::G1Element;
3use chia_streamable_macro::streamable;
4use chia_traits::chia_error;
5
6#[streamable(no_json)]
7pub struct ProofOfSpace {
8    challenge: Bytes32,
9    pool_public_key: Option<G1Element>,
10    pool_contract_puzzle_hash: Option<Bytes32>,
11    plot_public_key: G1Element,
12    /// The 2 top bits determine the type of proof:
13    /// 00 = v1 plot, the field store k-size
14    /// 10 = v2 plot, the field store strength
15    /// 01 = reserved
16    /// 11 = reserved
17    /// this field was renamed when adding support for v2 plots since the top
18    /// bit now means whether it's v1 or v2. To stay backwards compabible with
19    /// JSON serialization, we still serialize this as its original name
20    #[cfg_attr(feature = "serde", serde(rename = "size", alias = "version_and_size"))]
21    version_and_size: u8,
22    proof: Bytes,
23}
24
25/// The k-size for v1 PoS, or strength if it's a v2 PoS
26#[derive(Debug, PartialEq)]
27pub enum PlotParam {
28    KSize(u8),
29    Strength(u8),
30}
31
32impl ProofOfSpace {
33    pub fn param(&self) -> chia_error::Result<PlotParam> {
34        match self.version_and_size & 0b1100_0000 {
35            // valid v1 plot sizes are 32-50 (mainnet) and 18-50 (testnet)
36            0b0000_0000 => Ok(PlotParam::KSize(self.version_and_size)),
37            // valid v2 plot strength are 2-63
38            0b1000_0000 => Ok(PlotParam::Strength(self.version_and_size & 0x3f)),
39            _ => Err(chia_error::Error::InvalidPoSVersion),
40        }
41    }
42}
43
44#[cfg(feature = "py-bindings")]
45use chia_traits::{FromJsonDict, ToJsonDict};
46#[cfg(feature = "py-bindings")]
47use pyo3::prelude::*;
48
49#[cfg(feature = "py-bindings")]
50#[pyclass(name = "PlotParam")]
51pub struct PyPlotParam {
52    #[pyo3(get)]
53    pub size_v1: Option<u8>,
54    #[pyo3(get)]
55    pub strength_v2: Option<u8>,
56}
57
58#[cfg(feature = "py-bindings")]
59#[pymethods]
60impl PyPlotParam {
61    #[staticmethod]
62    fn make_v1(s: u8) -> Self {
63        assert!(s < 64);
64        Self {
65            size_v1: Some(s),
66            strength_v2: None,
67        }
68    }
69
70    #[staticmethod]
71    fn make_v2(s: u8) -> Self {
72        assert!(s < 64);
73        Self {
74            size_v1: None,
75            strength_v2: Some(s),
76        }
77    }
78}
79
80#[cfg(feature = "py-bindings")]
81#[pymethods]
82impl ProofOfSpace {
83    #[pyo3(name = "param")]
84    fn py_param(&self) -> PyResult<PyPlotParam> {
85        match self.param()? {
86            PlotParam::KSize(s) => Ok(PyPlotParam {
87                size_v1: Some(s),
88                strength_v2: None,
89            }),
90            PlotParam::Strength(s) => Ok(PyPlotParam {
91                size_v1: None,
92                strength_v2: Some(s),
93            }),
94        }
95    }
96}
97
98#[cfg(feature = "py-bindings")]
99impl ToJsonDict for ProofOfSpace {
100    fn to_json_dict(&self, py: pyo3::Python<'_>) -> pyo3::PyResult<pyo3::PyObject> {
101        use pyo3::prelude::PyDictMethods;
102        let ret = pyo3::types::PyDict::new(py);
103
104        ret.set_item("challenge", self.challenge.to_json_dict(py)?)?;
105        ret.set_item("pool_public_key", self.pool_public_key.to_json_dict(py)?)?;
106        ret.set_item(
107            "pool_contract_puzzle_hash",
108            self.pool_contract_puzzle_hash.to_json_dict(py)?,
109        )?;
110        ret.set_item("plot_public_key", self.plot_public_key.to_json_dict(py)?)?;
111
112        // "size" was the original name of this field. We keep it to remain backwards compatible
113        ret.set_item("size", self.version_and_size.to_json_dict(py)?)?;
114        ret.set_item("proof", self.proof.to_json_dict(py)?)?;
115
116        Ok(ret.into())
117    }
118}
119
120#[cfg(feature = "py-bindings")]
121impl FromJsonDict for ProofOfSpace {
122    fn from_json_dict(o: &pyo3::Bound<'_, pyo3::PyAny>) -> pyo3::PyResult<Self> {
123        use pyo3::prelude::PyAnyMethods;
124        Ok(Self {
125            challenge: <Bytes32 as FromJsonDict>::from_json_dict(&o.get_item("challenge")?)?,
126            pool_public_key: <Option<G1Element> as FromJsonDict>::from_json_dict(
127                &o.get_item("pool_public_key")?,
128            )?,
129            pool_contract_puzzle_hash: <Option<Bytes32> as FromJsonDict>::from_json_dict(
130                &o.get_item("pool_contract_puzzle_hash")?,
131            )?,
132            plot_public_key: <G1Element as FromJsonDict>::from_json_dict(
133                &o.get_item("plot_public_key")?,
134            )?,
135            version_and_size: <u8 as FromJsonDict>::from_json_dict(&o.get_item("size")?)?,
136            proof: <Bytes as FromJsonDict>::from_json_dict(&o.get_item("proof")?)?,
137        })
138    }
139}
140
141#[cfg(test)]
142#[allow(clippy::needless_pass_by_value)]
143mod tests {
144    use super::*;
145    use rstest::rstest;
146
147    #[rstest]
148    #[case(0x00, Ok(PlotParam::KSize(0)))]
149    #[case(0x01, Ok(PlotParam::KSize(1)))]
150    #[case(0x08, Ok(PlotParam::KSize(8)))]
151    #[case(0x3f, Ok(PlotParam::KSize(0x3f)))]
152    #[case(0x80, Ok(PlotParam::Strength(0)))]
153    #[case(0x81, Ok(PlotParam::Strength(1)))]
154    #[case(0x80 + 28, Ok(PlotParam::Strength(28)))]
155    #[case(0x80 + 30, Ok(PlotParam::Strength(30)))]
156    #[case(0x80 + 32, Ok(PlotParam::Strength(32)))]
157    #[case(0xff, Err(chia_error::Error::InvalidPoSVersion))]
158    #[case(0x7f, Err(chia_error::Error::InvalidPoSVersion))]
159    fn proof_of_space_size(#[case] size_field: u8, #[case] expect: chia_traits::Result<PlotParam>) {
160        let pos = ProofOfSpace::new(
161            Bytes32::from(b"abababababababababababababababab"),
162            None,
163            None,
164            G1Element::default(),
165            size_field,
166            Bytes::from(vec![]),
167        );
168
169        assert_eq!(pos.param(), expect);
170    }
171}