Skip to content

Commit

Permalink
feat(dre): [REL-1517] Auto-pick automation neuron for updating unassi…
Browse files Browse the repository at this point in the history
…gned nodes (#416)

* wip

* Fixed all known issues
  • Loading branch information
sasa-tomic committed May 23, 2024
1 parent 5bfb885 commit 20640cd
Show file tree
Hide file tree
Showing 8 changed files with 99 additions and 35 deletions.
81 changes: 58 additions & 23 deletions rs/cli/src/detect_neuron.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,24 +12,57 @@ use ic_nns_constants::GOVERNANCE_CANISTER_ID;
use ic_nns_governance::pb::v1::{ListNeurons, ListNeuronsResponse};
use ic_sys::utility_command::UtilityCommand;
use keyring::{Entry, Error};
use log::info;
use log::{info, warn};

#[derive(Clone)]
static RELEASE_AUTOMATION_DEFAULT_PRIVATE_KEY_PEM: &str = ".config/dfx/identity/release-automation/identity.pem"; // Relative to the home directory
const RELEASE_AUTOMATION_NEURON_ID: u64 = 80;

#[derive(Clone, Debug)]
pub struct Neuron {
network: ic_management_types::Network,
neuron_id: RefCell<Option<u64>>,
private_key_pem: Option<String>,
private_key_pem: Option<PathBuf>,
hsm_slot: Option<u64>,
hsm_pin: Option<String>,
hsm_key_id: Option<String>,
auth_cache: RefCell<Option<Auth>>,
}

impl Neuron {
pub async fn new(
network: &ic_management_types::Network,
neuron_id: Option<u64>,
private_key_pem: Option<String>,
hsm_slot: Option<u64>,
hsm_pin: Option<String>,
hsm_key_id: Option<String>,
) -> Self {
let private_key_pem = match private_key_pem {
Some(path) => match PathBuf::from_str(&path).expect("Cannot parse the private key path") {
path if path.exists() => Some(path),
_ => {
warn!("Invalid private key path");
None
}
},
None => None,
};
Self {
network: network.clone(),
neuron_id: RefCell::new(neuron_id),
private_key_pem,
hsm_slot,
hsm_pin,
hsm_key_id,
auth_cache: RefCell::new(None),
}
}

pub async fn get_auth(&self) -> anyhow::Result<Auth> {
if let Some(auth) = &*self.auth_cache.borrow() {
return Ok(auth.clone());
};

let auth = if let Some(path) = &self.private_key_pem {
Auth::Keyfile { path: path.clone() }
} else {
Expand Down Expand Up @@ -73,30 +106,26 @@ impl Neuron {
Ok(vec!["--proposer".to_string(), neuron_id.to_string()])
}

pub async fn new(
network: &ic_management_types::Network,
neuron_id: Option<u64>,
private_key_pem: Option<String>,
hsm_slot: Option<u64>,
hsm_pin: Option<String>,
hsm_key_id: Option<String>,
) -> Self {
pub fn as_automation(self) -> Self {
let private_key_pem = match self.private_key_pem {
Some(private_key_pem) => private_key_pem,
None => dirs::home_dir()
.expect("failed to find the home dir")
.join(RELEASE_AUTOMATION_DEFAULT_PRIVATE_KEY_PEM),
};

Self {
network: network.clone(),
neuron_id: RefCell::new(neuron_id),
private_key_pem,
hsm_slot,
hsm_pin,
hsm_key_id,
auth_cache: RefCell::new(None),
private_key_pem: Some(private_key_pem),
neuron_id: RefCell::new(Some(RELEASE_AUTOMATION_NEURON_ID)),
..self
}
}
}

#[derive(Clone)]
#[derive(Clone, Debug)]
pub enum Auth {
Hsm { pin: String, slot: u64, key_id: String },
Keyfile { path: String },
Keyfile { path: PathBuf },
}

fn pkcs11_lib_path() -> anyhow::Result<PathBuf> {
Expand Down Expand Up @@ -129,7 +158,7 @@ impl Auth {
"--key-id".to_string(),
key_id.clone(),
],
Auth::Keyfile { path } => vec!["--secret-key-pem".to_string(), path.clone()],
Auth::Keyfile { path } => vec!["--secret-key-pem".to_string(), path.to_string_lossy().to_string()],
}
}

Expand All @@ -140,7 +169,10 @@ impl Auth {
hsm_key_id: Option<String>,
) -> anyhow::Result<Self> {
match (private_key_pem, hsm_slot, hsm_pin, hsm_key_id) {
(Some(path), _, _, _) => Ok(Auth::Keyfile { path }),
(Some(path), _, _, _) if PathBuf::from(path.clone()).exists() => Ok(Auth::Keyfile {
path: PathBuf::from(path),
}),
(Some(path), _, _, _) => Err(anyhow::anyhow!("Invalid key file path: {}", path)),
(None, Some(slot), Some(pin), Some(key_id)) => Ok(Auth::Hsm { pin, slot, key_id }),
_ => Err(anyhow::anyhow!("Invalid auth arguments")),
}
Expand Down Expand Up @@ -210,7 +242,10 @@ pub async fn auto_detect_neuron_id(nns_urls: &[url::Url], auth: Auth) -> anyhow:
let response = Decode!(&response, ListNeuronsResponse)?;
let neuron_ids = response.neuron_infos.keys().copied().collect::<Vec<_>>();
match neuron_ids.len() {
0 => Err(anyhow::anyhow!("HSM doesn't control any neurons")),
0 => Err(anyhow::anyhow!(
"HSM doesn't control any neurons. Response fro governance canister: {:?}",
response
)),
1 => Ok(neuron_ids[0]),
_ => Select::with_theme(&ColorfulTheme::default())
.items(&neuron_ids)
Expand Down
23 changes: 17 additions & 6 deletions rs/cli/src/ic_admin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,15 @@ impl IcAdminWrapper {
}
}

pub fn as_automation(self) -> Self {
Self {
network: self.network,
ic_admin_bin_path: self.ic_admin_bin_path,
proceed_without_confirmation: self.proceed_without_confirmation,
neuron: self.neuron.as_automation(),
}
}

pub fn from_cli(cli: ParsedCli) -> Self {
Self {
network: cli.network,
Expand Down Expand Up @@ -203,6 +212,7 @@ impl IcAdminWrapper {
));
}
}

self.run(
&cmd.get_command_name(),
[
Expand All @@ -229,12 +239,12 @@ impl IcAdminWrapper {
]
})
.unwrap_or_default(),
self.neuron.as_arg_vec(!as_simulation).await?,
self.neuron.as_arg_vec(true).await?,
cmd.args(),
]
.concat()
.as_slice(),
!as_simulation,
true,
false,
)
.await
Expand All @@ -256,10 +266,11 @@ impl IcAdminWrapper {
self._exec(cmd.clone(), opts.clone(), true).await?;
}

if Confirm::new()
.with_prompt("Do you want to continue?")
.default(false)
.interact()?
if self.proceed_without_confirmation
|| Confirm::new()
.with_prompt("Do you want to continue?")
.default(false)
.interact()?
{
// User confirmed the desire to submit the proposal and no obvious problems were
// found. Proceeding!
Expand Down
5 changes: 5 additions & 0 deletions rs/cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,11 @@ async fn async_main() -> Result<(), anyhow::Error> {
},

cli::Commands::UpdateUnassignedNodes { nns_subnet_id } => {
let runner_instance = if target_network.is_mainnet() {
runner_instance.as_automation()
} else {
runner_instance
};
let nns_subnet_id = match nns_subnet_id {
Some(subnet_id) => subnet_id.to_owned(),
None => {
Expand Down
4 changes: 2 additions & 2 deletions rs/cli/src/registry_dump.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ fn get_node_rewards_table(
Ok(r) => match r {
Some(r) => r,
None => {
if network.name.eq("mainnet") {
if network.is_mainnet() {
panic!("Failed to get Node Rewards Table")
} else {
warn!("Didn't find any node rewards details for network: {}", network.name);
Expand All @@ -271,7 +271,7 @@ fn get_node_rewards_table(
}
},
Err(_) => {
if network.name.eq("mainnet") {
if network.is_mainnet() {
panic!("Failed to get Node Rewards Table for mainnet")
} else {
warn!("Failed to get Node Rewards Table for {}", network.name);
Expand Down
8 changes: 8 additions & 0 deletions rs/cli/src/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,14 @@ impl Runner {
Ok(())
}

pub fn as_automation(self) -> Self {
Self {
ic_admin: self.ic_admin.as_automation(),
dashboard_backend_client: self.dashboard_backend_client,
registry: self.registry,
}
}

pub async fn subnet_resize(
&self,
request: ic_management_types::requests::SubnetResizeRequest,
Expand Down
7 changes: 4 additions & 3 deletions rs/ic-management-backend/src/health.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,10 @@ enum HealthStatusQuerierImplementations {

impl From<Network> for HealthStatusQuerierImplementations {
fn from(value: Network) -> Self {
match value.name.as_str() {
"mainnet" => HealthStatusQuerierImplementations::Dashboard(PublicDashboardHealthClient::new(None)),
_ => HealthStatusQuerierImplementations::Prometheus(PrometheusHealthClient::new(value.clone())),
if value.is_mainnet() {
HealthStatusQuerierImplementations::Dashboard(PublicDashboardHealthClient::new(None))
} else {
HealthStatusQuerierImplementations::Prometheus(PrometheusHealthClient::new(value))
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion rs/ic-management-backend/src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ impl RegistryState {

pub fn update_node_labels_guests(&mut self, node_label_guests: Vec<Guest>) {
self.node_labels_guests = node_label_guests;
if self.network.name != "mainnet" {
if !self.network.is_mainnet() {
for g in &mut self.node_labels_guests {
g.dfinity_owned = true;
}
Expand Down
4 changes: 4 additions & 0 deletions rs/ic-management-types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -654,6 +654,10 @@ impl Network {
_ => self.name.clone(),
}
}

pub fn is_mainnet(&self) -> bool {
self.name == "mainnet"
}
}

impl std::fmt::Display for Network {
Expand Down

0 comments on commit 20640cd

Please sign in to comment.