replace expects/unwraps with fallible vrf routines

Author: kantai

stacks-common/src/util/vrf.rs

@@ -181,6 +181,7 @@ impl VRFPublicKey {
 pub enum Error {
     InvalidPublicKey,
     InvalidDataError,
+    InvalidHashPoints,
     OSRNGError(rand::Error),
 }
 
@@ -189,6 +190,7 @@ impl fmt::Display for Error {
         match *self {
             Error::InvalidPublicKey => write!(f, "Invalid public key"),
             Error::InvalidDataError => write!(f, "No data could be found"),
+            Error::InvalidHashPoints => write!(f, "VRF hash points did not yield a valid scalar"),
             Error::OSRNGError(ref e) => fmt::Display::fmt(e, f),
         }
     }
@@ -199,6 +201,7 @@ impl error::Error for Error {
         match *self {
             Error::InvalidPublicKey => None,
             Error::InvalidDataError => None,
+            Error::InvalidHashPoints => None,
             Error::OSRNGError(ref e) => Some(e),
         }
     }
@@ -474,17 +477,17 @@ impl VRF {
 
     /// Convert a 16-byte string into a scalar.
     /// The upper 16 bytes in the resulting scalar MUST BE 0's
-    fn ed25519_scalar_from_hash128(hash128: &[u8; 16]) -> ed25519_Scalar {
+    fn ed25519_scalar_from_hash128(hash128: &[u8; 16]) -> Option<ed25519_Scalar> {
         let mut scalar_buf = [0u8; 32];
         scalar_buf[0..16].copy_from_slice(hash128);
 
-        ed25519_Scalar::from_canonical_bytes(scalar_buf).expect("Invalid scalar")
+        ed25519_Scalar::from_canonical_bytes(scalar_buf).into()
     }
 
     /// ECVRF proof routine
     /// https://tools.ietf.org/id/draft-irtf-cfrg-vrf-02.html#rfc.section.5.1
     #[allow(clippy::op_ref)]
-    pub fn prove(secret: &VRFPrivateKey, alpha: &[u8]) -> VRFProof {
+    pub fn prove(secret: &VRFPrivateKey, alpha: &[u8]) -> Option<VRFProof> {
         let (Y_point, x_scalar, trunc_hash) = VRF::expand_privkey(secret);
         let H_point = VRF::hash_to_curve(&Y_point, alpha);
 
@@ -495,14 +498,15 @@ impl VRF {
         let kH_point = &k_scalar * &H_point;
 
         let c_hashbuf = VRF::hash_points(&H_point, &Gamma_point, &kB_point, &kH_point);
-        let c_scalar = VRF::ed25519_scalar_from_hash128(&c_hashbuf);
+        let c_scalar = VRF::ed25519_scalar_from_hash128(&c_hashbuf)?;
 
         let s_scalar = &k_scalar + &c_scalar * &x_scalar;
 
         // NOTE: expect() won't panic because c_scalar is guaranteed to have
         // its upper 16 bytes as 0
         VRFProof::new(Gamma_point, c_scalar, s_scalar)
-            .expect("FATAL ERROR: upper-16 bytes of proof's C scalar are NOT 0")
+            .inspect_err(|e| error!("FATAL: upper-16 bytes of proof's C scalar are NOT 0: {e}"))
+            .ok()
     }
 
     /// Given a public key, verify that the private key owner that generate the ECVRF proof did so on the given message.
@@ -525,7 +529,9 @@ impl VRF {
         let V_point = s_reduced * &H_point - proof.c() * proof.Gamma();
 
         let c_prime_hashbuf = VRF::hash_points(&H_point, proof.Gamma(), &U_point, &V_point);
-        let c_prime = VRF::ed25519_scalar_from_hash128(&c_prime_hashbuf);
+        let Some(c_prime) = VRF::ed25519_scalar_from_hash128(&c_prime_hashbuf) else {
+            return Err(Error::InvalidHashPoints);
+        };
 
         // NOTE: this leverages constant-time comparison inherited from the Scalar impl
         Ok(c_prime == *(proof.c()))
@@ -587,7 +593,7 @@ mod tests {
             let privk = VRFPrivateKey::from_bytes(&proof_fixture.privkey[..]).unwrap();
             let expected_proof_bytes = &proof_fixture.proof[..];
 
-            let proof = VRF::prove(&privk, &alpha.to_vec());
+            let proof = VRF::prove(&privk, &alpha.to_vec()).unwrap();
             let proof_bytes = proof.to_bytes();
 
             assert_eq!(proof_bytes.to_vec(), expected_proof_bytes.to_vec());
@@ -609,7 +615,7 @@ mod tests {
             let mut msg = [0u8; 1024];
             rng.fill_bytes(&mut msg);
 
-            let proof = VRF::prove(&secret_key, &msg);
+            let proof = VRF::prove(&secret_key, &msg).unwrap();
             let res = VRF::verify(&public_key, &proof, &msg).unwrap();
 
             assert!(res);

stackslib/src/burnchains/tests/mod.rs

@@ -241,7 +241,7 @@ impl TestMiner {
         );
         match self.vrf_key_map.get(vrf_pubkey) {
             Some(prover_key) => {
-                let proof = VRF::prove(prover_key, last_sortition_hash.as_bytes());
+                let proof = VRF::prove(prover_key, last_sortition_hash.as_bytes())?;
                 let valid = match VRF::verify(vrf_pubkey, &proof, last_sortition_hash.as_bytes()) {
                     Ok(v) => v,
                     Err(e) => false,

stackslib/src/chainstate/coordinator/tests.rs

@@ -646,7 +646,7 @@ fn make_genesis_block_with_recipients(
 
     let parent_stacks_header = StacksHeaderInfo::regtest_genesis();
 
-    let proof = VRF::prove(vrf_key, sortition_tip.sortition_hash.as_bytes());
+    let proof = VRF::prove(vrf_key, sortition_tip.sortition_hash.as_bytes()).unwrap();
 
     let mut builder = StacksBlockBuilder::make_regtest_block_builder(
         burnchain,
@@ -909,7 +909,7 @@ fn make_stacks_block_with_input(
 
     eprintln!("Build off of {:?}", &parent_stacks_header);
 
-    let proof = VRF::prove(vrf_key, sortition_tip.sortition_hash.as_bytes());
+    let proof = VRF::prove(vrf_key, sortition_tip.sortition_hash.as_bytes()).unwrap();
 
     let total_burn = parents_sortition.total_burn;
 

stackslib/src/chainstate/nakamoto/tests/mod.rs

@@ -1628,7 +1628,7 @@ fn test_nakamoto_block_static_verification() {
     let vrf_privkey = VRFPrivateKey::new();
     let vrf_pubkey = VRFPublicKey::from_private(&vrf_privkey);
     let sortition_hash = SortitionHash([0x01; 32]);
-    let vrf_proof = VRF::prove(&vrf_privkey, sortition_hash.as_bytes());
+    let vrf_proof = VRF::prove(&vrf_privkey, sortition_hash.as_bytes()).unwrap();
 
     let burn_recipient = StacksAddress::burn_address(false).to_account_principal();
     let alt_recipient = StacksAddress::p2pkh(false, &StacksPublicKey::from_private(&private_key_2))

testnet/stacks-node/src/keychain.rs

@@ -118,14 +118,38 @@ impl Keychain {
 
     /// Generate a VRF proof over a given byte message.
     /// `block_height` must be the _same_ block height called to make_vrf_keypair()
-    pub fn generate_proof(&self, block_height: u64, bytes: &[u8; 32]) -> VRFProof {
+    pub fn generate_proof(&self, block_height: u64, bytes: &[u8; 32]) -> Option<VRFProof> {
         let (pk, sk) = self.make_vrf_keypair(block_height);
-        let proof = VRF::prove(&sk, bytes.as_ref());
+        let Some(proof) = VRF::prove(&sk, bytes.as_ref()) else {
+            error!(
+                "Failed to generate proof with keypair, will be unable to mine.";
+                "block_height" => block_height,
+                "pk" => ?pk
+            );
+            return None;
+        };
 
         // Ensure that the proof is valid by verifying
-        let is_valid = VRF::verify(&pk, &proof, bytes.as_ref()).unwrap_or(false);
-        assert!(is_valid);
-        proof
+        let is_valid = VRF::verify(&pk, &proof, bytes.as_ref())
+            .inspect_err(|e| {
+                error!(
+                    "Failed to validate generated proof, will be unable to mine.";
+                    "block_height" => block_height,
+                    "pk" => ?pk,
+                    "err" => %e,
+                );
+            })
+            .ok()?;
+        if !is_valid {
+            error!(
+                "Generated invalidat proof, will be unable to mine.";
+                "block_height" => block_height,
+                "pk" => ?pk,
+            );
+            None
+        } else {
+            Some(proof)
+        }
     }
 
     /// Generate a microblock signing key for this burnchain block height.
@@ -367,7 +391,7 @@ mod tests {
             };
 
             // Generate the proof
-            let proof = VRF::prove(vrf_sk, bytes.as_ref());
+            let proof = VRF::prove(vrf_sk, bytes.as_ref())?;
             // Ensure that the proof is valid by verifying
             let is_valid = VRF::verify(vrf_pk, &proof, bytes.as_ref()).unwrap_or(false);
             assert!(is_valid);

testnet/stacks-node/src/nakamoto_node/miner.rs

@@ -1127,6 +1127,17 @@ impl BlockMinerThread {
             )
         };
 
+        let Some(vrf_proof) = vrf_proof else {
+            error!(
+                "Unable to generate VRF proof, will be unable to mine";
+                "burn_block_sortition_hash" => %self.burn_election_block.sortition_hash,
+                "burn_block_block_height" => %self.burn_block.block_height,
+                "burn_block_hash" => %self.burn_block.burn_header_hash,
+                "vrf_pubkey" => &self.registered_key.vrf_public_key.to_hex()
+            );
+            return None;
+        };
+
         debug!(
             "Generated VRF Proof: {} over {} ({},{}) with key {}",
             vrf_proof.to_hex(),

testnet/stacks-node/src/neon_node.rs

@@ -1730,6 +1730,17 @@ impl BlockMinerThread {
             )
         };
 
+        let Some(vrf_proof) = vrf_proof else {
+            error!(
+                "Unable to generate VRF proof, will be unable to mine";
+                "burn_block_sortition_hash" => %self.burn_block.sortition_hash,
+                "burn_block_block_height" => %self.burn_block.block_height,
+                "burn_block_hash" => %self.burn_block.burn_header_hash,
+                "vrf_pubkey" => &self.registered_key.vrf_public_key.to_hex()
+            );
+            return None;
+        };
+
         debug!(
             "Generated VRF Proof: {} over {} ({},{}) with key {}",
             vrf_proof.to_hex(),

testnet/stacks-node/src/node.rs

@@ -659,10 +659,13 @@ impl Node {
             .expect("FATAL: failed to query canonical burn chain tip");
 
         // Generates a proof out of the sortition hash provided in the params.
-        let vrf_proof = self.keychain.generate_proof(
+        let Some(vrf_proof) = self.keychain.generate_proof(
             registered_key.target_block_height,
             tip.sortition_hash.as_bytes(),
-        );
+        ) else {
+            warn!("Failed to generate VRF proof, will be unable to initiate new tenure");
+            return None;
+        };
 
         // Generates a new secret key for signing the trail of microblocks
         // of the upcoming tenure.
@@ -731,10 +734,13 @@ impl Node {
         if self.active_registered_key.is_some() {
             let registered_key = self.active_registered_key.clone().unwrap();
 
-            let vrf_proof = self.keychain.generate_proof(
+            let Some(vrf_proof) = self.keychain.generate_proof(
                 registered_key.target_block_height,
                 burnchain_tip.block_snapshot.sortition_hash.as_bytes(),
-            );
+            ) else {
+                warn!("Failed to generate VRF proof, will be unable to mine commits");
+                return;
+            };
 
             let op = self.generate_block_commit_op(
                 anchored_block_from_ongoing_tenure.header.block_hash(),