[safety-rules] Allow config based initialization of safety rules data

Sergvh · May 10, 2022 · 85cba9f · 85cba9f
1 parent a7e699d
commit 85cba9f
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 7 deletions.
diff --git a/config/src/config/mod.rs b/config/src/config/mod.rs
@@ -162,11 +162,11 @@ impl WaypointConfig {
 
 #[derive(Deserialize, Serialize)]
 pub struct IdentityBlob {
-    account_address: AccountAddress,
-    account_key: Ed25519PrivateKey,
+    pub account_address: AccountAddress,
+    pub account_key: Ed25519PrivateKey,
     /// Optional consensus key.  Only used for validators
-    consensus_key: Option<Ed25519PrivateKey>,
-    network_key: x25519::PrivateKey,
+    pub consensus_key: Option<Ed25519PrivateKey>,
+    pub network_key: x25519::PrivateKey,
 }
 
 impl IdentityBlob {

diff --git a/config/src/config/safety_rules_config.rs b/config/src/config/safety_rules_config.rs
@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    config::{LoggerConfig, SecureBackend},
+    config::{IdentityBlob, LoggerConfig, SecureBackend, WaypointConfig},
     keys::ConfigKey,
 };
 use aptos_crypto::{ed25519::Ed25519PrivateKey, Uniform};
@@ -26,6 +26,7 @@ pub struct SafetyRulesConfig {
     // Read/Write/Connect networking operation timeout in milliseconds.
     pub network_timeout_ms: u64,
     pub enable_cached_safety_data: bool,
+    pub initial_safety_rules_config: InitialSafetyRulesConfig,
 }
 
 impl Default for SafetyRulesConfig {
@@ -40,6 +41,7 @@ impl Default for SafetyRulesConfig {
             // Default value of 30 seconds for a timeout
             network_timeout_ms: 30_000,
             enable_cached_safety_data: true,
+            initial_safety_rules_config: InitialSafetyRulesConfig::None,
         }
     }
 }
@@ -52,6 +54,35 @@ impl SafetyRulesConfig {
     }
 }
 
+// TODO: Find a cleaner way so WaypointConfig isn't duplicated
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum InitialSafetyRulesConfig {
+    FromFile {
+        identity_blob_path: PathBuf,
+        waypoint: WaypointConfig,
+    },
+    None,
+}
+
+impl InitialSafetyRulesConfig {
+    pub fn waypoint(&self) -> Waypoint {
+        match self {
+            InitialSafetyRulesConfig::FromFile { waypoint, .. } => waypoint.waypoint(),
+            InitialSafetyRulesConfig::None => panic!("Must have a waypoint"),
+        }
+    }
+
+    pub fn identity_blob(&self) -> IdentityBlob {
+        match self {
+            InitialSafetyRulesConfig::FromFile {
+                identity_blob_path, ..
+            } => IdentityBlob::from_file(identity_blob_path).unwrap(),
+            InitialSafetyRulesConfig::None => panic!("Must have an identity blob"),
+        }
+    }
+}
+
 /// Defines how safety rules should be executed
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 #[serde(rename_all = "snake_case", tag = "type")]

diff --git a/consensus/safety-rules/src/safety_rules_manager.rs b/consensus/safety-rules/src/safety_rules_manager.rs
@@ -10,7 +10,7 @@ use crate::{
     thread::ThreadService,
     SafetyRules, TSafetyRules,
 };
-use aptos_config::config::{SafetyRulesConfig, SafetyRulesService};
+use aptos_config::config::{InitialSafetyRulesConfig, SafetyRulesConfig, SafetyRulesService};
 use aptos_infallible::RwLock;
 use aptos_secure_storage::{KVStorage, Storage};
 use std::{convert::TryInto, net::SocketAddr, sync::Arc};
@@ -45,7 +45,36 @@ pub fn storage(config: &SafetyRulesConfig) -> PersistentSafetyStorage {
             config.enable_cached_safety_data,
         )
     } else {
-        PersistentSafetyStorage::new(internal_storage, config.enable_cached_safety_data)
+        let storage =
+            PersistentSafetyStorage::new(internal_storage, config.enable_cached_safety_data);
+        // If it's initialized, then we can continue
+        if storage.author().is_ok() {
+            storage
+        } else if !matches!(
+            config.initial_safety_rules_config,
+            InitialSafetyRulesConfig::None
+        ) {
+            let identity_blob = config.initial_safety_rules_config.identity_blob();
+            let waypoint = config.initial_safety_rules_config.waypoint();
+
+            let backend = &config.backend;
+            let internal_storage: Storage =
+                backend.try_into().expect("Unable to initialize storage");
+            PersistentSafetyStorage::initialize(
+                internal_storage,
+                identity_blob.account_address,
+                identity_blob
+                    .consensus_key
+                    .expect("Consensus key needed for safety rules"),
+                identity_blob.account_key,
+                waypoint,
+                config.enable_cached_safety_data,
+            )
+        } else {
+            panic!(
+                "Safety rules storage is not initialized, provide an initial safety rules config"
+            )
+        }
     }
 }