chore(upgrader): introduce disaster recovery request operation

core/upgrader/api/spec.did

@@ -162,9 +162,9 @@ type SetDisasterRecoveryAccountsAndAssetsInput = record {
 
 // Request to trigger disaster recovery. Requests are stored in the Upgrader
 // canister, and when at least `quorum` of the committee members
-// agree on the exact module, args, and install mode, the request is processed.
+// agree on the exact disaster recovery input, the request is processed.
 // Requests older than 1 week will be discarded.
-type RequestDisasterRecoveryInput = record {
+type RequestDisasterRecoveryInstallCodeInput = record {
   // The wasm module to be installed.
   module : blob;
   // Additional wasm module chunks to append to the wasm module.
@@ -174,6 +174,9 @@ type RequestDisasterRecoveryInput = record {
   // The install mode: Install, Upgrade, or Reinstall.
   install_mode : InstallMode;
 };
+type RequestDisasterRecoveryInput = variant {
+  InstallCode : RequestDisasterRecoveryInstallCodeInput;
+};
 
 type InstallMode = variant {
   // Install the wasm module.
@@ -226,16 +229,25 @@ type GetLogsResult = variant {
   Err : Error;
 };
 
-// Request to recover the station.
-type StationRecoveryRequest = record {
-  // The requester user id.
-  user_id : text;
-  // The sha of the wasm module to be installed.
-  wasm_sha256 : blob;
+type StationRecoveryRequestInstallCodeOperation = record {
   // The install mode: Install, Upgrade, or Reinstall.
   install_mode : InstallMode;
+  // The sha of the wasm module to be installed.
+  wasm_sha256 : blob;
   // The argument to be passed to the install function.
   arg : blob;
+};
+
+type StationRecoveryRequestOperation = variant {
+  InstallCode : StationRecoveryRequestInstallCodeOperation;
+};
+
+// Request to recover the station.
+type StationRecoveryRequest = record {
+  // The requester user id.
+  user_id : text;
+  // The disaster recovery operation.
+  operation : StationRecoveryRequestOperation;
   // The request submission timestamp.
   submitted_at : text;
 };

core/upgrader/api/src/lib.rs

@@ -155,7 +155,7 @@ pub enum InstallMode {
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
-pub struct RequestDisasterRecoveryInput {
+pub struct RequestDisasterRecoveryInstallCodeInput {
     #[serde(with = "serde_bytes")]
     pub module: Vec<u8>,
     pub module_extra_chunks: Option<WasmModuleExtraChunks>,
@@ -165,6 +165,11 @@ pub struct RequestDisasterRecoveryInput {
     pub install_mode: InstallMode,
 }
 
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub enum RequestDisasterRecoveryInput {
+    InstallCode(RequestDisasterRecoveryInstallCodeInput),
+}
+
 #[derive(CandidType, Deserialize, Debug, Clone)]
 pub struct PaginationInput {
     pub offset: Option<u64>,
@@ -198,15 +203,26 @@ pub enum TriggerUpgradeResponse {
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
-pub struct StationRecoveryRequest {
-    /// The user ID of the station.
-    pub user_id: UuidDTO,
-    /// The SHA-256 hash of the wasm module.
-    pub wasm_sha256: Vec<u8>,
+pub struct StationRecoveryRequestInstallCodeOperation {
     /// The install mode: upgrade or reinstall.
     pub install_mode: InstallMode,
+    /// The SHA-256 hash of the wasm module.
+    pub wasm_sha256: Vec<u8>,
     /// The install arguments.
     pub arg: Vec<u8>,
+}
+
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub enum StationRecoveryRequestOperation {
+    InstallCode(StationRecoveryRequestInstallCodeOperation),
+}
+
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub struct StationRecoveryRequest {
+    /// The user ID of the station.
+    pub user_id: UuidDTO,
+    /// The disaster recovery operation.
+    pub operation: StationRecoveryRequestOperation,
     /// Time in nanoseconds since the UNIX epoch when the request was submitted.
     pub submitted_at: TimestampRfc3339,
 }

core/upgrader/impl/src/controllers/mod.rs

@@ -3,3 +3,18 @@ mod logs;
 
 pub use disaster_recovery::*;
 pub use logs::*;
+
+use crate::services::DISASTER_RECOVERY_SERVICE;
+use ic_cdk::post_upgrade;
+
+#[post_upgrade]
+fn post_upgrade() {
+    if !DISASTER_RECOVERY_SERVICE
+        .storage
+        .get()
+        .recovery_requests
+        .is_empty()
+    {
+        ic_cdk::trap("upgrader cannot be upgraded due to pending disaster recovery requests")
+    }
+}

core/upgrader/impl/src/lib.rs

@@ -21,6 +21,7 @@ pub use orbit_essentials::cdk::mocks as upgrader_ic_cdk;
 
 pub mod controllers;
 pub mod errors;
+pub mod mappers;
 pub mod model;
 pub mod services;
 pub mod upgrade;

core/upgrader/impl/src/mappers/disaster_recovery.rs

@@ -0,0 +1,62 @@
+use crate::model::{
+    RequestDisasterRecoveryInstallCodeLog, RequestDisasterRecoveryOperationLog,
+    StationRecoveryRequestInstallCodeOperation, StationRecoveryRequestOperation,
+};
+use orbit_essentials::utils::sha256_hash;
+
+impl From<&upgrader_api::RequestDisasterRecoveryInput> for StationRecoveryRequestOperation {
+    fn from(request: &upgrader_api::RequestDisasterRecoveryInput) -> Self {
+        match request {
+            upgrader_api::RequestDisasterRecoveryInput::InstallCode(install_code) => {
+                let wasm_sha256 =
+                    if let Some(ref module_extra_chunks) = install_code.module_extra_chunks {
+                        module_extra_chunks.wasm_module_hash.clone()
+                    } else {
+                        sha256_hash(&install_code.module)
+                    };
+                StationRecoveryRequestOperation::InstallCode(
+                    StationRecoveryRequestInstallCodeOperation {
+                        install_mode: install_code.install_mode.clone().into(),
+                        wasm_module: install_code.module.clone(),
+                        wasm_module_extra_chunks: install_code.module_extra_chunks.clone(),
+                        wasm_sha256,
+                        arg: install_code.arg.clone(),
+                        arg_sha256: sha256_hash(&install_code.arg),
+                    },
+                )
+            }
+        }
+    }
+}
+
+impl From<&StationRecoveryRequestOperation> for RequestDisasterRecoveryOperationLog {
+    fn from(operation: &StationRecoveryRequestOperation) -> Self {
+        match operation {
+            StationRecoveryRequestOperation::InstallCode(ref install_code) => {
+                RequestDisasterRecoveryOperationLog::InstallCode(
+                    RequestDisasterRecoveryInstallCodeLog {
+                        install_mode: install_code.install_mode.to_string(),
+                        wasm_sha256: hex::encode(&install_code.wasm_sha256),
+                        arg_sha256: hex::encode(&install_code.arg_sha256),
+                    },
+                )
+            }
+        }
+    }
+}
+
+impl From<&StationRecoveryRequestOperation> for upgrader_api::StationRecoveryRequestOperation {
+    fn from(operation: &StationRecoveryRequestOperation) -> Self {
+        match operation {
+            StationRecoveryRequestOperation::InstallCode(ref install_code) => {
+                upgrader_api::StationRecoveryRequestOperation::InstallCode(
+                    upgrader_api::StationRecoveryRequestInstallCodeOperation {
+                        install_mode: install_code.install_mode.into(),
+                        wasm_sha256: install_code.wasm_sha256.clone(),
+                        arg: install_code.arg.clone(),
+                    },
+                )
+            }
+        }
+    }
+}

core/upgrader/impl/src/mappers/mod.rs

@@ -0,0 +1 @@
+mod disaster_recovery;

core/upgrader/impl/src/model/disaster_recovery.rs

@@ -12,7 +12,7 @@ use uuid::Uuid;
 use crate::utils::HelperMapper;
 
 #[storable]
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 pub enum InstallMode {
     /// Install the wasm module.
     Install,
@@ -63,22 +63,35 @@ impl From<InstallMode> for CanisterInstallMode {
 }
 
 #[storable]
-#[derive(Clone, Debug)]
-pub struct StationRecoveryRequest {
-    /// The user ID of the station.
-    pub user_id: UUID,
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct StationRecoveryRequestInstallCodeOperation {
+    /// The install mode: upgrade or reinstall.
+    pub install_mode: InstallMode,
     /// The wasm module to be installed.
     pub wasm_module: Vec<u8>,
     /// Optional extra chunks of the wasm module to be installed.
     pub wasm_module_extra_chunks: Option<WasmModuleExtraChunks>,
     /// The SHA-256 hash of the wasm module.
     pub wasm_sha256: Vec<u8>,
-    /// The install mode: upgrade or reinstall.
-    pub install_mode: InstallMode,
     /// The install arguments.
     pub arg: Vec<u8>,
     /// The SHA-256 hash of the install arguments.
     pub arg_sha256: Vec<u8>,
+}
+
+#[storable]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub enum StationRecoveryRequestOperation {
+    InstallCode(StationRecoveryRequestInstallCodeOperation),
+}
+
+#[storable]
+#[derive(Clone, Debug)]
+pub struct StationRecoveryRequest {
+    /// The user ID of the station.
+    pub user_id: UUID,
+    /// The disaster recovery operation.
+    pub operation: StationRecoveryRequestOperation,
     /// Time in nanoseconds since the UNIX epoch when the request was submitted.
     pub submitted_at: Timestamp,
 }
@@ -87,9 +100,7 @@ impl From<StationRecoveryRequest> for upgrader_api::StationRecoveryRequest {
     fn from(value: StationRecoveryRequest) -> Self {
         upgrader_api::StationRecoveryRequest {
             user_id: Uuid::from_bytes(value.user_id).hyphenated().to_string(),
-            wasm_sha256: value.wasm_sha256,
-            install_mode: upgrader_api::InstallMode::from(value.install_mode),
-            arg: value.arg,
+            operation: (&value.operation).into(),
             submitted_at: timestamp_to_rfc3339(&value.submitted_at),
         }
     }

core/upgrader/impl/src/model/logging.rs

@@ -29,18 +29,40 @@ pub struct SetAccountsAndAssetsLog {
 }
 
 #[derive(Serialize)]
-pub struct RequestDisasterRecoveryLog {
-    pub user: AdminUser,
+pub struct RequestDisasterRecoveryInstallCodeLog {
+    pub install_mode: String,
     pub wasm_sha256: String,
     pub arg_sha256: String,
-    pub install_mode: String,
+}
+
+#[derive(Serialize)]
+pub enum RequestDisasterRecoveryOperationLog {
+    InstallCode(RequestDisasterRecoveryInstallCodeLog),
+}
+
+impl std::fmt::Display for RequestDisasterRecoveryOperationLog {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            RequestDisasterRecoveryOperationLog::InstallCode(install_code) => {
+                write!(
+                    f,
+                    "InstallCode with mode{}, wasm hash {}, and arg hash {}",
+                    install_code.install_mode, install_code.wasm_sha256, install_code.arg_sha256
+                )
+            }
+        }
+    }
+}
+
+#[derive(Serialize)]
+pub struct RequestDisasterRecoveryLog {
+    pub user: AdminUser,
+    pub operation: RequestDisasterRecoveryOperationLog,
 }
 
 #[derive(Serialize)]
 pub struct DisasterRecoveryStartLog {
-    pub wasm_sha256: String,
-    pub arg_sha256: String,
-    pub install_mode: String,
+    pub operation: RequestDisasterRecoveryOperationLog,
 }
 
 #[derive(Serialize)]
@@ -109,16 +131,14 @@ impl LogEntryType {
                 format!("Set {} disaster recovery account(s)", data.accounts.len(),)
             }
             LogEntryType::RequestDisasterRecovery(data) => format!(
-                "{} requested disaster recovery with wasm hash {} and arg hash {}",
+                "{} requested disaster recovery with operation {}",
                 data.user.to_summary(),
-                hex::encode(&data.wasm_sha256),
-                hex::encode(&data.arg_sha256)
+                data.operation,
             ),
 
             LogEntryType::DisasterRecoveryStart(data) => format!(
-                "Disaster recovery successfully initiated to {} station with wasm {}",
-                data.install_mode,
-                hex::encode(&data.wasm_sha256)
+                "Disaster recovery successfully initiated with operation {}",
+                data.operation,
             ),
             LogEntryType::DisasterRecoveryResult(data) => match data.result {
                 RecoveryResult::Success => "Disaster recovery succeeded".to_owned(),