feat: Adding possibility to override neurons per command (#1159)

dfinity · Dec 19, 2024 · a2e006d · a2e006d
1 parent 4e20800
commit a2e006d
Show file tree

Hide file tree

Showing 14 changed files with 130 additions and 191 deletions.
diff --git a/rs/cli/src/auth.rs b/rs/cli/src/auth.rs
@@ -95,6 +95,7 @@ impl Neuron {
         network: &Network,
         neuron_id: Option<u64>,
         offline: bool,
+        neuron_override: Option<Neuron>,
     ) -> anyhow::Result<Self> {
         let (neuron_id, auth_opts) = if network.name == "staging" {
             let staging_known_path = dirs::home_dir().expect("Home dir should be set").join(STAGING_KEY_PATH_FROM_HOME);
@@ -128,42 +129,66 @@ impl Neuron {
             (neuron_id, auth_opts)
         };
 
+        let auth_specified = !matches!(
+            auth_opts,
+            AuthOpts {
+                private_key_pem: None,
+                hsm_opts: HsmOpts {
+                    hsm_pin: None,
+                    hsm_params: HsmParams {
+                        hsm_slot: None,
+                        hsm_key_id: None,
+                    },
+                },
+            }
+        );
+
         match requirement {
             AuthRequirement::Anonymous => Ok(Self {
                 auth: Auth::Anonymous,
                 neuron_id: 0,
                 include_proposer: false,
             }),
             AuthRequirement::Signer => Ok(Self {
-                auth: Auth::from_auth_opts(auth_opts).await?,
+                // If nothing is specified for the signer and override is provided
+                // use overide neuron for auth
+                auth: match neuron_override {
+                    Some(neuron) if !auth_specified => neuron.auth,
+                    _ => Auth::from_auth_opts(auth_opts).await?,
+                },
                 neuron_id: 0,
                 include_proposer: false,
             }),
             AuthRequirement::Neuron => Ok({
-                match (neuron_id, offline) {
-                    (Some(n), _) => Self {
-                        neuron_id: n,
-                        auth: Auth::from_auth_opts(auth_opts).await?,
-                        include_proposer: true,
-                    },
-                    // This is just a placeholder since
-                    // the tool is instructed to run in
-                    // offline mode.
-                    (None, true) => {
-                        warn!("Required full neuron but offline mode instructed! Will not attempt to auto-detect neuron id");
-                        Self {
-                            neuron_id: 0,
+                if neuron_id.is_none() && !auth_specified && neuron_override.is_some() {
+                    info!("Using override neuron for this command since no auth options were provided");
+                    neuron_override.unwrap()
+                } else {
+                    match (neuron_id, offline) {
+                        (Some(n), _) => Self {
+                            neuron_id: n,
                             auth: Auth::from_auth_opts(auth_opts).await?,
                             include_proposer: true,
+                        },
+                        // This is just a placeholder since
+                        // the tool is instructed to run in
+                        // offline mode.
+                        (None, true) => {
+                            warn!("Required full neuron but offline mode instructed! Will not attempt to auto-detect neuron id");
+                            Self {
+                                neuron_id: 0,
+                                auth: Auth::from_auth_opts(auth_opts).await?,
+                                include_proposer: true,
+                            }
                         }
-                    }
-                    (None, false) => {
-                        let auth = Auth::from_auth_opts(auth_opts).await?;
-                        let neuron_id = auth.clone().auto_detect_neuron_id(network.nns_urls.clone()).await?;
-                        Self {
-                            neuron_id,
-                            auth,
-                            include_proposer: true,
+                        (None, false) => {
+                            let auth = Auth::from_auth_opts(auth_opts).await?;
+                            let neuron_id = auth.clone().auto_detect_neuron_id(network.nns_urls.clone()).await?;
+                            Self {
+                                neuron_id,
+                                auth,
+                                include_proposer: true,
+                            }
                         }
                     }
                 }
@@ -228,6 +253,12 @@ impl Neuron {
     }
 }
 
+const AUTOMATION_NEURON_DEFAULT_PATH: &str = ".config/dfx/identity/release-automation/identity.pem";
+pub fn get_automation_neuron_default_path() -> PathBuf {
+    let home = dirs::home_dir().unwrap();
+    home.join(AUTOMATION_NEURON_DEFAULT_PATH)
+}
+
 #[derive(Debug, Clone)]
 pub enum Auth {
     Hsm { identity: ParallelHardwareIdentity },

diff --git a/rs/cli/src/commands/api_boundary_nodes/mod.rs b/rs/cli/src/commands/api_boundary_nodes/mod.rs
@@ -12,21 +12,7 @@ mod update;
 #[derive(Args, Debug)]
 pub struct ApiBoundaryNodes {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-impl_executable_command_for_enums! { Add, Update, Remove }
-
-impl ExecutableCommand for ApiBoundaryNodes {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
-}
+impl_executable_command_for_enums! { ApiBoundaryNodes, Add, Update, Remove }
diff --git a/rs/cli/src/commands/hostos/mod.rs b/rs/cli/src/commands/hostos/mod.rs
@@ -10,21 +10,7 @@ pub mod rollout_from_node_group;
 #[derive(Args, Debug)]
 pub struct HostOs {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-super::impl_executable_command_for_enums! { Rollout, RolloutFromNodeGroup }
-
-impl ExecutableCommand for HostOs {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
-}
+super::impl_executable_command_for_enums! { HostOs, Rollout, RolloutFromNodeGroup }
diff --git a/rs/cli/src/commands/mod.rs b/rs/cli/src/commands/mod.rs
@@ -220,30 +220,8 @@ The argument is mandatory for testnets, and is optional for mainnet and staging"
     pub cordoned_features_file: Option<String>,
 }
 
-// Do not use outside of DRE CLI.
-// You can run your command by directly instantiating it.
-impl ExecutableCommand for Args {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommands.require_auth()
-    }
-
-    async fn execute(&self, ctx: DreContext) -> anyhow::Result<()> {
-        self.subcommands.execute(ctx).await
-    }
-
-    /// Validate the command line arguments. You can return an error with something like:
-    /// ```rust
-    /// if args.neuron_id.is_none() {
-    ///    cmd.error(ErrorKind::MissingRequiredArgument, "Neuron ID is required for this command.")).exit();
-    /// }
-    /// ```
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut Command) {
-        self.subcommands.validate(args, cmd)
-    }
-}
-
 macro_rules! impl_executable_command_for_enums {
-    ($($var:ident),*) => {
+    ($str_name:ident, $($var:ident),*) => {
         use crate::ctx::DreContext;
         use clap::{Subcommand, Command};
 
@@ -270,19 +248,53 @@ macro_rules! impl_executable_command_for_enums {
                     $(Subcommands::$var(variant) => variant.validate(args, cmd),)*
                 }
             }
+
+            fn neuron_override(&self) -> Option<crate::auth::Neuron> {
+                match &self {
+                    $(Subcommands::$var(variant) => variant.neuron_override(),)*
+                }
+            }
+        }
+
+        impl ExecutableCommand for $str_name {
+            fn require_auth(&self) -> AuthRequirement {
+                self.subcommands.require_auth()
+            }
+
+            async fn execute(&self, ctx: DreContext) -> anyhow::Result<()> {
+                self.subcommands.execute(ctx).await
+            }
+
+            /// Validate the command line arguments. You can return an error with something like:
+            /// ```rust
+            /// if args.neuron_id.is_none() {
+            ///    cmd.error(ErrorKind::MissingRequiredArgument, "Neuron ID is required for this command.")).exit();
+            /// }
+            /// ```
+            fn validate(&self, args: &crate::commands::Args, cmd: &mut Command) {
+                self.subcommands.validate(args, cmd)
+            }
+
+            fn neuron_override(&self) -> Option<crate::auth::Neuron> {
+                self.subcommands.neuron_override()
+            }
         }
     }
 }
 pub(crate) use impl_executable_command_for_enums;
 
-impl_executable_command_for_enums! { DerToPrincipal, Network, Subnet, Get, Propose, UpdateUnassignedNodes, Version, NodeMetrics, HostOs, Nodes, ApiBoundaryNodes, Vote, Registry, Firewall, Upgrade, Proposals, Completions, Qualify, UpdateAuthorizedSubnets, Neuron }
+impl_executable_command_for_enums! { Args, DerToPrincipal, Network, Subnet, Get, Propose, UpdateUnassignedNodes, Version, NodeMetrics, HostOs, Nodes, ApiBoundaryNodes, Vote, Registry, Firewall, Upgrade, Proposals, Completions, Qualify, UpdateAuthorizedSubnets, Neuron }
 
 pub trait ExecutableCommand {
     fn require_auth(&self) -> AuthRequirement;
 
     fn validate(&self, args: &crate::commands::Args, cmd: &mut Command);
 
     fn execute(&self, ctx: DreContext) -> impl std::future::Future<Output = anyhow::Result<()>>;
+
+    fn neuron_override(&self) -> Option<crate::auth::Neuron> {
+        None
+    }
 }
 
 #[derive(Clone)]

diff --git a/rs/cli/src/commands/neuron/mod.rs b/rs/cli/src/commands/neuron/mod.rs
@@ -12,21 +12,7 @@ mod top_up;
 #[derive(Args, Debug)]
 pub struct Neuron {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-impl_executable_command_for_enums! { Balance, TopUp, Refresh }
-
-impl ExecutableCommand for Neuron {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut Command) {
-        self.subcommand.validate(args, cmd)
-    }
-
-    async fn execute(&self, ctx: DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-}
+impl_executable_command_for_enums! { Neuron, Balance, TopUp, Refresh }
diff --git a/rs/cli/src/commands/nodes/mod.rs b/rs/cli/src/commands/nodes/mod.rs
@@ -8,20 +8,6 @@ mod remove;
 #[derive(Args, Debug)]
 pub struct Nodes {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
-}
-impl_executable_command_for_enums! { Remove }
-
-impl ExecutableCommand for Nodes {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
+    pub subcommands: Subcommands,
 }
+impl_executable_command_for_enums! { Nodes, Remove }
diff --git a/rs/cli/src/commands/proposals/mod.rs b/rs/cli/src/commands/proposals/mod.rs
@@ -59,24 +59,10 @@ mod pending;
 #[clap(alias = "proposal")]
 pub struct Proposals {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-impl_executable_command_for_enums! { Pending, Get, Analyze, Filter, List }
-
-impl ExecutableCommand for Proposals {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
-}
+impl_executable_command_for_enums! { Proposals, Pending, Get, Analyze, Filter, List }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Proposal {

diff --git a/rs/cli/src/commands/qualify/mod.rs b/rs/cli/src/commands/qualify/mod.rs
@@ -10,21 +10,7 @@ mod list;
 #[derive(Args, Debug)]
 pub struct Qualify {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-impl_executable_command_for_enums! { List, Execute }
-
-impl ExecutableCommand for Qualify {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-}
+impl_executable_command_for_enums! { Qualify, List, Execute }
diff --git a/rs/cli/src/commands/subnet/deploy.rs b/rs/cli/src/commands/subnet/deploy.rs
@@ -2,9 +2,13 @@ use clap::Args;
 
 use ic_types::PrincipalId;
 
-use crate::commands::{AuthRequirement, ExecutableCommand};
+use crate::{
+    auth::get_automation_neuron_default_path,
+    commands::{AuthRequirement, ExecutableCommand},
+};
 
 #[derive(Args, Debug)]
+#[clap(visible_aliases = &["upgrade", "update"])]
 pub struct Deploy {
     /// Version to propose for the subnet
     #[clap(long, short)]
@@ -29,4 +33,14 @@ impl ExecutableCommand for Deploy {
     }
 
     fn validate(&self, _args: &crate::commands::Args, _cmd: &mut clap::Command) {}
+
+    fn neuron_override(&self) -> Option<crate::auth::Neuron> {
+        Some(crate::auth::Neuron {
+            auth: crate::auth::Auth::Keyfile {
+                path: get_automation_neuron_default_path(),
+            },
+            neuron_id: 80,
+            include_proposer: true,
+        })
+    }
 }
diff --git a/rs/cli/src/commands/subnet/mod.rs b/rs/cli/src/commands/subnet/mod.rs
@@ -18,21 +18,7 @@ mod whatif;
 #[derive(Parser, Debug)]
 pub struct Subnet {
     #[clap(subcommand)]
-    pub subcommand: Subcommands,
+    pub subcommands: Subcommands,
 }
 
-impl_executable_command_for_enums! { WhatifDecentralization, Deploy, Replace, Resize, Create, Rescue }
-
-impl ExecutableCommand for Subnet {
-    fn require_auth(&self) -> AuthRequirement {
-        self.subcommand.require_auth()
-    }
-
-    async fn execute(&self, ctx: crate::ctx::DreContext) -> anyhow::Result<()> {
-        self.subcommand.execute(ctx).await
-    }
-
-    fn validate(&self, args: &crate::commands::Args, cmd: &mut clap::Command) {
-        self.subcommand.validate(args, cmd)
-    }
-}
+impl_executable_command_for_enums! { Subnet, WhatifDecentralization, Deploy, Replace, Resize, Create, Rescue }