fix PyVrfSk

dojoengine · Dec 10, 2023 · 44f89bf · 44f89bf
1 parent fcc24ed
commit 44f89bf
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/crates/native_blockifier/src/py_block_executor.rs b/crates/native_blockifier/src/py_block_executor.rs
@@ -2,14 +2,14 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use blockifier::block_context::BlockContext;
-use ecvrf::{VrfPk, VrfSk};
+use ecvrf::VrfPk;
 use pyo3::prelude::*;
 use starknet_api::block::{BlockNumber, BlockTimestamp};
 use starknet_api::core::{ChainId, ContractAddress};
 
 use crate::errors::NativeBlockifierResult;
 use crate::py_transaction_executor::PyTransactionExecutor;
-use crate::py_utils::{int_to_chain_id, py_attr, PyFelt};
+use crate::py_utils::{int_to_chain_id, py_attr, PyFelt, PyVrfSk};
 use crate::storage::Storage;
 
 #[pyclass]
@@ -53,7 +53,7 @@ pub struct PyGeneralConfig {
     pub cairo_resource_fee_weights: Arc<HashMap<String, f64>>,
     pub invoke_tx_max_n_steps: u32,
     pub validate_max_n_steps: u32,
-    pub ecvrf_private_key: PyVrfPk,
+    pub ecvrf_private_key: PyVrfSk,
 }
 
 impl FromPyObject<'_> for PyGeneralConfig {

diff --git a/crates/native_blockifier/src/py_utils.rs b/crates/native_blockifier/src/py_utils.rs
@@ -1,7 +1,7 @@
 use std::convert::TryFrom;
 
 use blockifier::transaction::errors::TransactionExecutionError;
-use ecvrf::VrfPk;
+use ecvrf::VrfSk;
 use num_bigint::BigUint;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
@@ -15,7 +15,7 @@ use crate::errors::NativeBlockifierResult;
 pub struct PyFelt(#[pyo3(from_py_with = "int_to_stark_felt")] pub StarkFelt);
 
 #[derive(Eq, FromPyObject, Hash, PartialEq, Clone, Copy)]
-pub struct PyVrfPk(#[pyo3(from_py_with = "int_to_stark_felt")] pub VrfPk);
+pub struct PyVrfSk(#[pyo3(from_py_with = "int_to_stark_felt")] pub VrfSk);
 
 impl IntoPy<PyObject> for PyFelt {
     fn into_py(self, py: Python<'_>) -> PyObject {
@@ -68,6 +68,18 @@ fn int_to_stark_felt(int: &PyAny) -> PyResult<StarkFelt> {
     biguint_to_felt(biguint).map_err(|e| PyValueError::new_err(e.to_string()))
 }
 
+fn int_to_vrf_sk(int: &PyAny) -> PyResult<VrfSk> {
+    let biguint: BigUint = int.extract()?;
+    let bytes = biguint.to_bytes_le();
+
+    // Create a fixed-size array and copy the bytes into it.
+    let mut sk_bytes = [0u8; 32];
+    let len = bytes.len().min(sk_bytes.len());
+    sk_bytes[..len].copy_from_slice(&bytes[..len]);
+
+    Ok(VrfSk::from(sk_bytes))
+}
+
 // TODO: Convert to a `TryFrom` cast and put in starknet-api (In StarkFelt).
 pub fn biguint_to_felt(biguint: BigUint) -> NativeBlockifierResult<StarkFelt> {
     let biguint_hex = format!("{biguint:#x}");