aura to babe keystore migration

Author: liamaharon

node/src/service.rs

@@ -9,13 +9,17 @@ use sc_consensus::{BasicQueue, BoxBlockImport};
 use sc_consensus_grandpa::BlockNumberOps;
 use sc_consensus_slots::BackoffAuthoringOnFinalizedHeadLagging;
 use sc_consensus_slots::SlotProportion;
+use sc_keystore::LocalKeystore;
 use sc_network::config::SyncMode;
 use sc_network_sync::strategy::warp::{WarpSyncConfig, WarpSyncProvider};
 use sc_service::{Configuration, PartialComponents, TaskManager, error::Error as ServiceError};
 use sc_telemetry::{Telemetry, TelemetryHandle, TelemetryWorker, log};
 use sc_transaction_pool::TransactionPoolHandle;
 use sc_transaction_pool_api::OffchainTransactionPoolFactory;
 use sp_core::H256;
+use sp_core::crypto::KeyTypeId;
+use sp_keystore::Keystore;
+use sp_runtime::key_types;
 use sp_runtime::traits::{Block as BlockT, NumberFor};
 use std::collections::HashSet;
 use std::str::FromStr;
@@ -31,6 +35,8 @@ use crate::ethereum::{
     StorageOverrideHandler, db_config_dir, new_frontier_partial, spawn_frontier_tasks,
 };
 
+const LOG_TARGET: &str = "node-service";
+
 /// The minimum period of blocks on which justifications will be
 /// imported and generated.
 const GRANDPA_JUSTIFICATION_PERIOD: u32 = 512;
@@ -94,6 +100,13 @@ pub fn new_partial(
             executor,
         )?;
 
+    // Prepare keystore for authoring Babe blocks.
+    copy_keys(
+        &keystore_container.local_keystore(),
+        key_types::AURA,
+        key_types::BABE,
+    )?;
+
     let client = Arc::new(client);
 
     let telemetry = telemetry.map(|(worker, telemetry)| {
@@ -758,3 +771,47 @@ fn run_manual_seal_authorship(
         .spawn_blocking("manual-seal", None, manual_seal);
     Ok(())
 }
+
+/// Copy `from_key_type` keys to also exist as `to_key_type`.
+///
+/// Used for the Aura to Babe migration, where Aura validators need their keystore to copy their
+/// Aura keys over to Babe. This works because Aura and Babe keys use identical crypto.
+fn copy_keys(
+    keystore: &LocalKeystore,
+    from_key_type: KeyTypeId,
+    to_key_type: KeyTypeId,
+) -> sc_keystore::Result<()> {
+    use std::collections::HashSet;
+
+    let from_keys: HashSet<_> = keystore
+        .raw_public_keys(from_key_type)?
+        .into_iter()
+        .collect();
+    let to_keys: HashSet<_> = keystore.raw_public_keys(to_key_type)?.into_iter().collect();
+    let to_copy: Vec<_> = from_keys.difference(&to_keys).collect();
+
+    log::debug!(target: LOG_TARGET, "from_keys: {:?}", from_keys);
+    log::debug!(target: LOG_TARGET, "to_keys: {:?}", to_keys);
+    log::debug!(target: LOG_TARGET, "to_copy: {:?} from {:?} to {:?}", &to_copy, from_key_type, to_key_type);
+
+    for public in to_copy {
+        if let Some(phrase) = keystore.key_phrase_by_type(&public, from_key_type)? {
+            if let Err(_) = keystore.insert(to_key_type, &phrase, &public) {
+                log::error!(
+                    target: LOG_TARGET,
+                    "Failed to copy key {:?} into keystore, insert operation failed.",
+                    &public,
+                );
+            };
+        } else {
+            log::error!(
+                target: LOG_TARGET,
+                "Failed to copy key from {:?} to {:?} as the key phrase is not available",
+                from_key_type,
+                to_key_type
+            );
+        }
+    }
+
+    Ok(())
+}