feat(upgrader): disaster recovery flow reuses the same upgrader canis…

…ter (#293) Enables the disaster recovery flow to reuse the same upgrader canister in case the mode of recovery would be set to `install` or `reinstall`.
dfinity · Jul 4, 2024 · a2fcdb5 · a2fcdb5
1 parent 9497822
commit a2fcdb5
Show file tree

Hide file tree

Showing 10 changed files with 237 additions and 37 deletions.
diff --git a/apps/wallet/src/generated/station/station.did b/apps/wallet/src/generated/station/station.did
@@ -1986,6 +1986,14 @@ type AdminInitInput = record {
   identity : principal;
 };
 
+// An input type for configuring the upgrader canister.
+type SystemUpgraderInput = variant {
+  // An existing upgrader canister.
+  Id : principal;
+  // Creates and deploys a new canister.
+  WasmModule : blob;
+};
+
 // The init configuration for the canister.
 //
 // Only used when installing the canister for the first time.
@@ -1996,8 +2004,8 @@ type SystemInit = record {
   admins : vec AdminInitInput;
   // Quorum of admins for initial policies.
   quorum : opt nat16;
-  // The wasm module of the station upgrader canister.
-  upgrader_wasm_module : blob;
+  // The upgrader configuration.
+  upgrader : SystemUpgraderInput;
   // An optional additional controller of the station and upgrader canisters.
   fallback_controller : opt principal;
 };

diff --git a/apps/wallet/src/generated/station/station.did.d.ts b/apps/wallet/src/generated/station/station.did.d.ts
@@ -865,8 +865,8 @@ export type SystemInfoResult = { 'Ok' : { 'system' : SystemInfo } } |
 export interface SystemInit {
   'name' : string,
   'fallback_controller' : [] | [Principal],
+  'upgrader' : SystemUpgraderInput,
   'admins' : Array<AdminInitInput>,
-  'upgrader_wasm_module' : Uint8Array | number[],
   'quorum' : [] | [number],
 }
 export type SystemInstall = { 'Upgrade' : SystemUpgrade } |
@@ -875,6 +875,8 @@ export type SystemResourceAction = { 'ManageSystemInfo' : null } |
   { 'SystemInfo' : null } |
   { 'Capabilities' : null };
 export interface SystemUpgrade { 'name' : [] | [string] }
+export type SystemUpgraderInput = { 'Id' : Principal } |
+  { 'WasmModule' : Uint8Array | number[] };
 export type TimestampRFC3339 = string;
 export interface Transfer {
   'id' : UUID,

diff --git a/apps/wallet/src/generated/station/station.did.js b/apps/wallet/src/generated/station/station.did.js
@@ -2,15 +2,19 @@ export const idlFactory = ({ IDL }) => {
   const RequestPolicyRule = IDL.Rec();
   const RequestPolicyRuleResult = IDL.Rec();
   const SystemUpgrade = IDL.Record({ 'name' : IDL.Opt(IDL.Text) });
+  const SystemUpgraderInput = IDL.Variant({
+    'Id' : IDL.Principal,
+    'WasmModule' : IDL.Vec(IDL.Nat8),
+  });
   const AdminInitInput = IDL.Record({
     'name' : IDL.Text,
     'identity' : IDL.Principal,
   });
   const SystemInit = IDL.Record({
     'name' : IDL.Text,
     'fallback_controller' : IDL.Opt(IDL.Principal),
+    'upgrader' : SystemUpgraderInput,
     'admins' : IDL.Vec(AdminInitInput),
-    'upgrader_wasm_module' : IDL.Vec(IDL.Nat8),
     'quorum' : IDL.Opt(IDL.Nat16),
   });
   const SystemInstall = IDL.Variant({
@@ -1183,15 +1187,19 @@ export const idlFactory = ({ IDL }) => {
 };
 export const init = ({ IDL }) => {
   const SystemUpgrade = IDL.Record({ 'name' : IDL.Opt(IDL.Text) });
+  const SystemUpgraderInput = IDL.Variant({
+    'Id' : IDL.Principal,
+    'WasmModule' : IDL.Vec(IDL.Nat8),
+  });
   const AdminInitInput = IDL.Record({
     'name' : IDL.Text,
     'identity' : IDL.Principal,
   });
   const SystemInit = IDL.Record({
     'name' : IDL.Text,
     'fallback_controller' : IDL.Opt(IDL.Principal),
+    'upgrader' : SystemUpgraderInput,
     'admins' : IDL.Vec(AdminInitInput),
-    'upgrader_wasm_module' : IDL.Vec(IDL.Nat8),
     'quorum' : IDL.Opt(IDL.Nat16),
   });
   const SystemInstall = IDL.Variant({

diff --git a/core/control-panel/impl/src/services/deploy.rs b/core/control-panel/impl/src/services/deploy.rs
@@ -104,7 +104,7 @@ impl DeployService {
             arg: Encode!(&station_api::SystemInstall::Init(station_api::SystemInit {
                 name: input.name.clone(),
                 admins,
-                upgrader_wasm_module,
+                upgrader: station_api::SystemUpgraderInput::WasmModule(upgrader_wasm_module),
                 quorum: Some(1),
                 fallback_controller: Some(NNS_ROOT_CANISTER_ID),
             }))

diff --git a/core/station/api/spec.did b/core/station/api/spec.did
@@ -1986,6 +1986,14 @@ type AdminInitInput = record {
   identity : principal;
 };
 
+// An input type for configuring the upgrader canister.
+type SystemUpgraderInput = variant {
+  // An existing upgrader canister.
+  Id : principal;
+  // Creates and deploys a new canister.
+  WasmModule : blob;
+};
+
 // The init configuration for the canister.
 //
 // Only used when installing the canister for the first time.
@@ -1996,8 +2004,8 @@ type SystemInit = record {
   admins : vec AdminInitInput;
   // Quorum of admins for initial policies.
   quorum : opt nat16;
-  // The wasm module of the station upgrader canister.
-  upgrader_wasm_module : blob;
+  // The upgrader configuration.
+  upgrader : SystemUpgraderInput;
   // An optional additional controller of the station and upgrader canisters.
   fallback_controller : opt principal;
 };

diff --git a/core/station/api/src/system.rs b/core/station/api/src/system.rs
@@ -32,6 +32,12 @@ pub struct AdminInitInput {
     pub identity: Principal,
 }
 
+#[derive(CandidType, serde::Serialize, Deserialize, Clone, Debug)]
+pub enum SystemUpgraderInput {
+    Id(Principal),
+    WasmModule(#[serde(with = "serde_bytes")] Vec<u8>),
+}
+
 #[derive(CandidType, serde::Serialize, Deserialize, Clone, Debug)]
 pub struct SystemInit {
     /// The station name.
@@ -40,9 +46,8 @@ pub struct SystemInit {
     pub admins: Vec<AdminInitInput>,
     /// The quorum of admin approvals required in initial policies.
     pub quorum: Option<u16>,
-    /// The upgrader canister wasm module.
-    #[serde(with = "serde_bytes")]
-    pub upgrader_wasm_module: Vec<u8>,
+    /// The upgrader configuration.
+    pub upgrader: SystemUpgraderInput,
     pub fallback_controller: Option<Principal>,
 }
 

diff --git a/core/station/impl/src/services/system.rs b/core/station/impl/src/services/system.rs
@@ -153,17 +153,15 @@ impl SystemService {
             print("Adding initial canister configurations");
             install_canister_handlers::init_post_process(&init).await?;
 
-            print("Deploying upgrader canister");
+            print("Init upgrader canister");
             let canister_id = self_canister_id();
             let mut upgrader_controllers = vec![canister_id];
             if let Some(fallback_controller) = init.fallback_controller {
                 upgrader_controllers.push(fallback_controller);
             }
-            let upgrader_canister_id = install_canister_handlers::deploy_upgrader(
-                init.upgrader_wasm_module,
-                upgrader_controllers,
-            )
-            .await?;
+            let upgrader_canister_id =
+                install_canister_handlers::init_upgrader(init.upgrader, upgrader_controllers)
+                    .await?;
             system_info.set_upgrader_canister_id(upgrader_canister_id);
 
             // sets the upgrader as a controller of the station canister
@@ -413,8 +411,32 @@ mod install_canister_handlers {
         Ok(())
     }
 
+    pub async fn init_upgrader(
+        input: station_api::SystemUpgraderInput,
+        controllers: Vec<Principal>,
+    ) -> Result<Principal, String> {
+        match input {
+            station_api::SystemUpgraderInput::Id(upgrader_id) => {
+                mgmt::update_settings(mgmt::UpdateSettingsArgument {
+                    canister_id: upgrader_id,
+                    settings: mgmt::CanisterSettings {
+                        controllers: Some(controllers),
+                        ..Default::default()
+                    },
+                })
+                .await
+                .map_err(|e| format!("Failed to set upgrader controller: {:?}", e))?;
+
+                Ok(upgrader_id)
+            }
+            station_api::SystemUpgraderInput::WasmModule(upgrader_wasm_module) => {
+                deploy_upgrader(upgrader_wasm_module, controllers).await
+            }
+        }
+    }
+
     /// Deploys the station upgrader canister and sets the station as the controller of the upgrader.
-    pub async fn deploy_upgrader(
+    async fn deploy_upgrader(
         upgrader_wasm_module: Vec<u8>,
         controllers: Vec<Principal>,
     ) -> Result<Principal, String> {
@@ -505,7 +527,7 @@ mod tests {
                     identity: Principal::from_slice(&[1; 29]),
                 }],
                 quorum: Some(1),
-                upgrader_wasm_module: vec![],
+                upgrader: station_api::SystemUpgraderInput::WasmModule(vec![]),
                 fallback_controller: None,
             })
             .await;

diff --git a/core/upgrader/impl/src/services/install_canister.rs b/core/upgrader/impl/src/services/install_canister.rs
@@ -1,11 +1,9 @@
-use std::sync::Arc;
-
+use crate::model::InstallMode;
 use async_trait::async_trait;
 use candid::Principal;
 use ic_cdk::api::management_canister::main::{self as mgmt, CanisterIdRecord, InstallCodeArgument};
 use lazy_static::lazy_static;
-
-use crate::model::InstallMode;
+use std::{collections::BTreeSet, sync::Arc};
 
 lazy_static! {
     pub static ref INSTALL_CANISTER: Arc<StationDisasterRecoveryInstall> =
@@ -39,6 +37,48 @@ impl InstallCanister for StationDisasterRecoveryInstall {
         arg: Vec<u8>,
         mode: InstallMode,
     ) -> Result<(), String> {
+        // For install and reinstall, we need to make the station self controlled.
+        match mode {
+            InstallMode::Install | InstallMode::Reinstall => {
+                // Get the current controllers and add the station to it. We preserve the existing controllers
+                // during this step to avoid losing control in case of a failed install.
+                let (info,) = mgmt::canister_info(mgmt::CanisterInfoRequest {
+                    canister_id,
+                    num_requested_changes: None,
+                })
+                .await
+                .map_err(|(code, err)| {
+                    format!(
+                        "failed to get canister info for canister: \"{}\", rejection code: {}",
+                        err, code as i32
+                    )
+                })?;
+
+                let mut controllers = BTreeSet::new();
+                controllers.extend(info.controllers.iter().cloned());
+                controllers.insert(canister_id);
+                controllers.insert(ic_cdk::id());
+
+                mgmt::update_settings(mgmt::UpdateSettingsArgument {
+                    canister_id,
+                    settings: mgmt::CanisterSettings {
+                        controllers: Some(controllers.into_iter().collect()),
+                        ..Default::default()
+                    },
+                })
+                .await
+                .map_err(|(code, err)| {
+                    format!(
+                        "failed to update settings for canister: \"{}\", rejection code: {}",
+                        err, code as i32
+                    )
+                })?;
+            }
+            InstallMode::Upgrade => {
+                // For upgrade, there are no controller changes needed.
+            }
+        }
+
         mgmt::install_code(InstallCodeArgument {
             mode: mode.into(),
             canister_id,