Merge pull request #4996 from stacks-network/feat/reuse-key

feat: support resuming from a saved VRF key

Author: obycode

testnet/stacks-node/src/nakamoto_node.rs

@@ -14,9 +14,10 @@
 // You should have received a copy of the GNU General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 use std::collections::HashSet;
+use std::io::Write;
 use std::sync::mpsc::Receiver;
-use std::thread;
 use std::thread::JoinHandle;
+use std::{fs, thread};
 
 use stacks::burnchains::{BurnchainSigner, Txid};
 use stacks::chainstate::burn::db::sortdb::SortitionDB;
@@ -277,6 +278,7 @@ impl StacksNode {
     /// Called from the main thread.
     pub fn process_burnchain_state(
         &mut self,
+        config: &Config,
         sortdb: &SortitionDB,
         sort_id: &SortitionId,
         ibd: bool,
@@ -316,9 +318,18 @@ impl StacksNode {
 
         let num_key_registers = key_registers.len();
 
-        self.globals
+        let activated_key_opt = self
+            .globals
             .try_activate_leader_key_registration(block_height, key_registers);
 
+        // save the registered VRF key
+        if let (Some(activated_key), Some(path)) = (
+            activated_key_opt,
+            config.miner.activated_vrf_key_path.as_ref(),
+        ) {
+            save_activated_vrf_key(path, &activated_key);
+        }
+
         debug!(
             "Processed burnchain state";
             "burn_height" => block_height,
@@ -339,3 +350,27 @@ impl StacksNode {
         self.p2p_thread_handle.join().unwrap();
     }
 }
+
+pub(crate) fn save_activated_vrf_key(path: &str, activated_key: &RegisteredKey) {
+    info!("Activated VRF key; saving to {}", path);
+
+    let Ok(key_json) = serde_json::to_string(&activated_key) else {
+        warn!("Failed to serialize VRF key");
+        return;
+    };
+
+    let mut f = match fs::File::create(&path) {
+        Ok(f) => f,
+        Err(e) => {
+            warn!("Failed to create {}: {:?}", &path, &e);
+            return;
+        }
+    };
+
+    if let Err(e) = f.write_all(key_json.as_str().as_bytes()) {
+        warn!("Failed to write activated VRF key to {}: {:?}", &path, &e);
+        return;
+    }
+
+    info!("Saved activated VRF key to {}", &path);
+}

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

@@ -15,6 +15,8 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 use core::fmt;
 use std::collections::HashSet;
+use std::fs;
+use std::io::Read;
 use std::sync::mpsc::{Receiver, RecvTimeoutError};
 use std::thread::JoinHandle;
 use std::time::{Duration, Instant};
@@ -1095,6 +1097,43 @@ impl RelayerThread {
         debug!("Relayer exit!");
     }
 
+    /// Try loading up a saved VRF key
+    pub(crate) fn load_saved_vrf_key(path: &str, pubkey_hash: &Hash160) -> Option<RegisteredKey> {
+        let mut f = match fs::File::open(path) {
+            Ok(f) => f,
+            Err(e) => {
+                warn!("Could not open {}: {:?}", &path, &e);
+                return None;
+            }
+        };
+        let mut registered_key_bytes = vec![];
+        if let Err(e) = f.read_to_end(&mut registered_key_bytes) {
+            warn!(
+                "Failed to read registered key bytes from {}: {:?}",
+                path, &e
+            );
+            return None;
+        }
+
+        let Ok(registered_key) = serde_json::from_slice::<RegisteredKey>(&registered_key_bytes)
+        else {
+            warn!(
+                "Did not load registered key from {}: could not decode JSON",
+                &path
+            );
+            return None;
+        };
+
+        // Check that the loaded key's memo matches the current miner's key
+        if registered_key.memo != pubkey_hash.as_ref() {
+            warn!("Loaded VRF key does not match mining key");
+            return None;
+        }
+
+        info!("Loaded registered key from {}", &path);
+        Some(registered_key)
+    }
+
     /// Top-level dispatcher
     pub fn handle_directive(&mut self, directive: RelayerDirective) -> bool {
         debug!("Relayer: handling directive"; "directive" => %directive);
@@ -1113,7 +1152,18 @@ impl RelayerThread {
                     info!("In initial block download, will not submit VRF registration");
                     return true;
                 }
-                self.rotate_vrf_and_register(&last_burn_block);
+                let mut saved_key_opt = None;
+                if let Some(path) = self.config.miner.activated_vrf_key_path.as_ref() {
+                    saved_key_opt =
+                        Self::load_saved_vrf_key(&path, &self.keychain.get_nakamoto_pkh());
+                }
+                if let Some(saved_key) = saved_key_opt {
+                    debug!("Relayer: resuming VRF key");
+                    self.globals.resume_leader_key(saved_key);
+                } else {
+                    self.rotate_vrf_and_register(&last_burn_block);
+                    debug!("Relayer: directive Registered VRF key");
+                }
                 self.globals.counters.bump_blocks_processed();
                 true
             }
@@ -1154,3 +1204,121 @@ impl RelayerThread {
         continue_running
     }
 }
+
+#[cfg(test)]
+pub mod test {
+    use std::fs::File;
+    use std::io::Write;
+    use std::path::Path;
+
+    use stacks::util::hash::Hash160;
+    use stacks::util::secp256k1::Secp256k1PublicKey;
+    use stacks::util::vrf::VRFPublicKey;
+
+    use super::RelayerThread;
+    use crate::nakamoto_node::save_activated_vrf_key;
+    use crate::run_loop::RegisteredKey;
+    use crate::Keychain;
+
+    #[test]
+    fn load_nonexistent_vrf_key() {
+        let keychain = Keychain::default(vec![0u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+
+        let path = "/tmp/does_not_exist.json";
+        _ = std::fs::remove_file(&path);
+
+        let res = RelayerThread::load_saved_vrf_key(&path, &pubkey_hash);
+        assert!(res.is_none());
+    }
+
+    #[test]
+    fn load_empty_vrf_key() {
+        let keychain = Keychain::default(vec![0u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+
+        let path = "/tmp/empty.json";
+        File::create(&path).expect("Failed to create test file");
+        assert!(Path::new(&path).exists());
+
+        let res = RelayerThread::load_saved_vrf_key(&path, &pubkey_hash);
+        assert!(res.is_none());
+
+        std::fs::remove_file(&path).expect("Failed to delete test file");
+    }
+
+    #[test]
+    fn load_bad_vrf_key() {
+        let keychain = Keychain::default(vec![0u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+
+        let path = "/tmp/invalid_saved_key.json";
+        let json_content = r#"{ "hello": "world" }"#;
+
+        // Write the JSON content to the file
+        let mut file = File::create(&path).expect("Failed to create test file");
+        file.write_all(json_content.as_bytes())
+            .expect("Failed to write to test file");
+        assert!(Path::new(&path).exists());
+
+        let res = RelayerThread::load_saved_vrf_key(&path, &pubkey_hash);
+        assert!(res.is_none());
+
+        std::fs::remove_file(&path).expect("Failed to delete test file");
+    }
+
+    #[test]
+    fn save_load_vrf_key() {
+        let keychain = Keychain::default(vec![0u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+        let key = RegisteredKey {
+            target_block_height: 101,
+            block_height: 102,
+            op_vtxindex: 1,
+            vrf_public_key: VRFPublicKey::from_hex(
+                "1da75863a7e1ef86f0f550d92b1f77dc60af23694b884b2816b703137ff94e71",
+            )
+            .unwrap(),
+            memo: pubkey_hash.as_ref().to_vec(),
+        };
+        let path = "/tmp/vrf_key.json";
+        save_activated_vrf_key(path, &key);
+
+        let res = RelayerThread::load_saved_vrf_key(&path, &pubkey_hash);
+        assert!(res.is_some());
+
+        std::fs::remove_file(&path).expect("Failed to delete test file");
+    }
+
+    #[test]
+    fn invalid_saved_memo() {
+        let keychain = Keychain::default(vec![0u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+        let key = RegisteredKey {
+            target_block_height: 101,
+            block_height: 102,
+            op_vtxindex: 1,
+            vrf_public_key: VRFPublicKey::from_hex(
+                "1da75863a7e1ef86f0f550d92b1f77dc60af23694b884b2816b703137ff94e71",
+            )
+            .unwrap(),
+            memo: pubkey_hash.as_ref().to_vec(),
+        };
+        let path = "/tmp/vrf_key.json";
+        save_activated_vrf_key(path, &key);
+
+        let keychain = Keychain::default(vec![1u8; 32]);
+        let pk = Secp256k1PublicKey::from_private(keychain.get_nakamoto_sk());
+        let pubkey_hash = Hash160::from_node_public_key(&pk);
+
+        let res = RelayerThread::load_saved_vrf_key(&path, &pubkey_hash);
+        assert!(res.is_none());
+
+        std::fs::remove_file(&path).expect("Failed to delete test file");
+    }
+}

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

@@ -653,7 +653,7 @@ impl SignCoordinator {
 
         let block_proposal_message = SignerMessageV0::BlockProposal(block_proposal);
         debug!("Sending block proposal message to signers";
-            "signer_signature_hash" => ?&block.header.signer_signature_hash().0,
+            "signer_signature_hash" => %block.header.signer_signature_hash(),
         );
         Self::send_miners_message_scalar::<SignerMessageV0>(
             &self.message_key,

testnet/stacks-node/src/neon_node.rs

@@ -3971,7 +3971,6 @@ impl RelayerThread {
                 if let Some(saved_key) = saved_key_opt {
                     self.globals.resume_leader_key(saved_key);
                 } else {
-                    debug!("Relayer: directive Register VRF key");
                     self.rotate_vrf_and_register(&last_burn_block);
                     debug!("Relayer: directive Registered VRF key");
                 }

testnet/stacks-node/src/run_loop/nakamoto.rs

@@ -635,9 +635,12 @@ impl RunLoop {
                         let sortition_id = &block.sortition_id;
 
                         // Have the node process the new block, that can include, or not, a sortition.
-                        if let Err(e) =
-                            node.process_burnchain_state(burnchain.sortdb_mut(), sortition_id, ibd)
-                        {
+                        if let Err(e) = node.process_burnchain_state(
+                            self.config(),
+                            burnchain.sortdb_mut(),
+                            sortition_id,
+                            ibd,
+                        ) {
                             // relayer errored, exit.
                             error!("Runloop: Block relayer and miner errored, exiting."; "err" => ?e);
                             return;