Ensure StorageIndex is only controller before installing StorageBucket (

#6070)
open-chat-labs · Jul 19, 2024 · 7b13a3e · 7b13a3e
1 parent 8ef8e90
commit 7b13a3e
Show file tree

Hide file tree

Showing 10 changed files with 71 additions and 81 deletions.
diff --git a/backend/canisters/group_index/CHANGELOG.md b/backend/canisters/group_index/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Clear old data from the failed upgrades log ([#6062](https://github.com/open-chat-labs/open-chat/pull/6062))
+- Ensure GroupIndex is only controller before installing LocalGroupIndex ([#6070](https://github.com/open-chat-labs/open-chat/pull/6070))
 
 ## [[2.0.1240](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1240-group_index)] - 2024-07-17
 

diff --git a/backend/canisters/group_index/impl/src/updates/add_local_group_index_canister.rs b/backend/canisters/group_index/impl/src/updates/add_local_group_index_canister.rs
@@ -3,49 +3,41 @@ use crate::{mutate_state, read_state, RuntimeState};
 use canister_api_macros::proposal;
 use canister_tracing_macros::trace;
 use group_index_canister::add_local_group_index_canister::{Response::*, *};
-use ic_cdk::api::management_canister::main::CanisterInstallMode;
 use tracing::info;
 use types::{BuildVersion, CanisterId, CanisterWasm};
-use utils::canister::{install, CanisterToInstall, WasmToInstall};
+use utils::canister::{install_basic, set_controllers};
 
 #[proposal(guard = "caller_is_governance_principal")]
 #[trace]
 async fn add_local_group_index_canister(args: Args) -> Response {
     match read_state(|state| prepare(&args, state)) {
         Ok(result) => {
             let wasm_version = result.canister_wasm.version;
-            match install(CanisterToInstall {
-                canister_id: args.canister_id,
-                current_wasm_version: BuildVersion::default(),
-                new_wasm_version: result.canister_wasm.version,
-                new_wasm: WasmToInstall::Default(result.canister_wasm.module),
-                deposit_cycles_if_needed: true,
-                args: result.init_args,
-                mode: CanisterInstallMode::Install,
-                stop_start_canister: false,
-            })
-            .await
-            {
-                Ok(_) => {
-                    let response = mutate_state(|state| commit(args.canister_id, wasm_version, state));
-                    info!(canister_id = %args.canister_id, "local group index canister added");
-                    response
-                }
-                Err(error) => InternalError(format!("{error:?}")),
+
+            if let Err(error) = set_controllers(args.canister_id, vec![result.this_canister_id]).await {
+                InternalError(format!("Failed to set controller: {error:?}"))
+            } else if let Err(error) = install_basic(args.canister_id, result.canister_wasm, result.init_args).await {
+                InternalError(format!("Failed to install canister: {error:?}"))
+            } else {
+                let response = mutate_state(|state| commit(args.canister_id, wasm_version, state));
+                info!(canister_id = %args.canister_id, "local group index canister added");
+                response
             }
         }
         Err(response) => response,
     }
 }
 
 struct PrepareResult {
+    this_canister_id: CanisterId,
     canister_wasm: CanisterWasm,
     init_args: local_group_index_canister::init::Args,
 }
 
 fn prepare(args: &Args, state: &RuntimeState) -> Result<PrepareResult, Response> {
     if !state.data.local_index_map.contains_key(&args.canister_id) {
         Ok(PrepareResult {
+            this_canister_id: state.env.canister_id(),
             canister_wasm: state.data.local_group_index_canister_wasm_for_new_canisters.clone(),
             init_args: local_group_index_canister::init::Args {
                 group_canister_wasm: state.data.group_canister_wasm.clone(),

diff --git a/backend/canisters/notifications_index/CHANGELOG.md b/backend/canisters/notifications_index/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Clear old data from the failed upgrades log ([#6062](https://github.com/open-chat-labs/open-chat/pull/6062))
+- Ensure NotificationsIndex is only controller before installing a Notifications canister ([#6070](https://github.com/open-chat-labs/open-chat/pull/6070))
 
 ## [[2.0.1219](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1219-notifications_index)] - 2024-07-03
 

diff --git a/backend/canisters/notifications_index/impl/src/updates/add_notifications_canister.rs b/backend/canisters/notifications_index/impl/src/updates/add_notifications_canister.rs
@@ -2,47 +2,41 @@ use crate::guards::caller_is_governance_principal;
 use crate::{mutate_state, read_state, NotificationsCanister, RuntimeState};
 use canister_api_macros::proposal;
 use canister_tracing_macros::trace;
-use ic_cdk::api::management_canister::main::CanisterInstallMode;
 use notifications_index_canister::add_notifications_canister::{Response::*, *};
 use notifications_index_canister::{NotificationsIndexEvent, SubscriptionAdded};
 use std::collections::hash_map::Entry::Vacant;
 use types::{BuildVersion, CanisterId, CanisterWasm};
-use utils::canister::{install, CanisterToInstall, WasmToInstall};
+use utils::canister::{install_basic, set_controllers};
 
 #[proposal(guard = "caller_is_governance_principal")]
 #[trace]
 async fn add_notifications_canister(args: Args) -> Response {
     match read_state(|state| prepare(args.canister_id, args.authorizers, state)) {
         Ok(result) => {
             let wasm_version = result.canister_wasm.version;
-            match install(CanisterToInstall {
-                canister_id: args.canister_id,
-                current_wasm_version: BuildVersion::default(),
-                new_wasm_version: result.canister_wasm.version,
-                new_wasm: WasmToInstall::Default(result.canister_wasm.module),
-                deposit_cycles_if_needed: true,
-                args: result.init_args,
-                mode: CanisterInstallMode::Install,
-                stop_start_canister: false,
-            })
-            .await
-            {
-                Ok(_) => mutate_state(|state| commit(args.canister_id, wasm_version, state)),
-                Err(error) => InternalError(format!("{error:?}")),
+
+            if let Err(error) = set_controllers(args.canister_id, vec![result.this_canister_id]).await {
+                InternalError(format!("Failed to set controller: {error:?}"))
+            } else if let Err(error) = install_basic(args.canister_id, result.canister_wasm, result.init_args).await {
+                InternalError(format!("Failed to install canister: {error:?}"))
+            } else {
+                mutate_state(|state| commit(args.canister_id, wasm_version, state))
             }
         }
         Err(response) => response,
     }
 }
 
 struct PrepareResult {
+    this_canister_id: CanisterId,
     canister_wasm: CanisterWasm,
     init_args: notifications_canister::init::Args,
 }
 
 fn prepare(canister_id: CanisterId, authorizers: Vec<CanisterId>, state: &RuntimeState) -> Result<PrepareResult, Response> {
     if !state.data.notifications_canisters.contains_key(&canister_id) {
         Ok(PrepareResult {
+            this_canister_id: state.env.canister_id(),
             canister_wasm: state.data.notifications_canister_wasm_for_new_canisters.clone(),
             init_args: notifications_canister::init::Args {
                 notifications_index_canister_id: state.env.canister_id(),

diff --git a/backend/canisters/storage_index/CHANGELOG.md b/backend/canisters/storage_index/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Clear old data from the failed upgrades log ([#6062](https://github.com/open-chat-labs/open-chat/pull/6062))
+- Ensure StorageIndex is only controller before installing StorageBucket ([#6070](https://github.com/open-chat-labs/open-chat/pull/6070))
 
 ## [[2.0.1175](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1175-storage_index)] - 2024-05-16
 

diff --git a/backend/canisters/storage_index/impl/src/updates/add_bucket_canister.rs b/backend/canisters/storage_index/impl/src/updates/add_bucket_canister.rs
@@ -5,22 +5,28 @@ use crate::{mutate_state, RuntimeState};
 use canister_api_macros::proposal;
 use canister_tracing_macros::trace;
 use storage_index_canister::add_bucket_canister::{Response::*, *};
-use types::{CanisterId, CanisterWasm, Cycles};
-use utils::canister::create_and_install;
+use types::{CanisterId, CanisterWasm};
+use utils::canister::{install_basic, set_controllers};
 
 // dfx canister --network ic call storage_index add_bucket_canister '(record { canister_id = principal "myzmx-wqaaa-aaaar-ad2ua-cai" })'
 #[proposal(guard = "caller_is_governance_principal")]
 #[trace]
 async fn add_bucket_canister(args: Args) -> Response {
-    let InitBucketArgs { wasm, init_args } = match read_state(|state| prepare(args.canister_id, state)) {
+    let InitBucketArgs {
+        this_canister_id,
+        wasm,
+        init_args,
+    } = match read_state(|state| prepare(args.canister_id, state)) {
         Ok(ok) => ok,
         Err(response) => return response,
     };
 
     let wasm_version = wasm.version;
 
-    if let Err(error) = create_and_install(Some(args.canister_id), wasm, init_args, 0, on_bucket_created).await {
-        InternalError(format!("{error:?}"))
+    if let Err(error) = set_controllers(args.canister_id, vec![this_canister_id]).await {
+        InternalError(format!("Failed to set controller: {error:?}"))
+    } else if let Err(error) = install_basic(args.canister_id, wasm, init_args).await {
+        InternalError(format!("Failed to install canister: {error:?}"))
     } else {
         let bucket = BucketRecord::new(args.canister_id, wasm_version);
         mutate_state(|state| state.data.add_bucket(bucket, false));
@@ -29,6 +35,7 @@ async fn add_bucket_canister(args: Args) -> Response {
 }
 
 struct InitBucketArgs {
+    this_canister_id: CanisterId,
     wasm: CanisterWasm,
     init_args: storage_bucket_canister::init::Args,
 }
@@ -38,6 +45,7 @@ fn prepare(canister_id: CanisterId, state: &RuntimeState) -> Result<InitBucketAr
         Err(BucketAlreadyAdded)
     } else {
         Ok(InitBucketArgs {
+            this_canister_id: state.env.canister_id(),
             wasm: state.data.bucket_canister_wasm.clone(),
             init_args: storage_bucket_canister::init::Args {
                 wasm_version: state.data.bucket_canister_wasm.version,
@@ -46,7 +54,3 @@ fn prepare(canister_id: CanisterId, state: &RuntimeState) -> Result<InitBucketAr
         })
     }
 }
-
-fn on_bucket_created(cycles: Cycles) {
-    mutate_state(|state| state.data.total_cycles_spent_on_canisters += cycles);
-}
diff --git a/backend/canisters/user_index/CHANGELOG.md b/backend/canisters/user_index/CHANGELOG.md
@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Fix fee then retry transfer if fee too high ([#6063](https://github.com/open-chat-labs/open-chat/pull/6063))
 - Handle transfer fee changing in either direction ([#6064](https://github.com/open-chat-labs/open-chat/pull/6064))
 - Accept proofs of uniqueness from LocalUserIndexes ([#6068](https://github.com/open-chat-labs/open-chat/pull/6068))
+- Ensure UserIndex is only controller before installing LocalUserIndex ([#6070](https://github.com/open-chat-labs/open-chat/pull/6070))
 
 ## [[2.0.1242](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1242-user_index)] - 2024-07-17
 

diff --git a/backend/canisters/user_index/impl/src/updates/add_local_user_index_canister.rs b/backend/canisters/user_index/impl/src/updates/add_local_user_index_canister.rs
@@ -2,51 +2,43 @@ use crate::guards::caller_is_governance_principal;
 use crate::{mutate_state, read_state, RuntimeState};
 use canister_api_macros::proposal;
 use canister_tracing_macros::trace;
-use ic_cdk::api::management_canister::main::CanisterInstallMode;
 use local_user_index_canister::{Event, UserRegistered};
 use tracing::info;
 use types::{BuildVersion, CanisterId, CanisterWasm};
 use user_index_canister::add_local_user_index_canister::{Response::*, *};
-use utils::canister::{install, CanisterToInstall, WasmToInstall};
+use utils::canister::{install_basic, set_controllers};
 
 #[proposal(guard = "caller_is_governance_principal")]
 #[trace]
 async fn add_local_user_index_canister(args: Args) -> Response {
     match read_state(|state| prepare(&args, state)) {
         Ok(result) => {
             let wasm_version = result.canister_wasm.version;
-            match install(CanisterToInstall {
-                canister_id: args.canister_id,
-                current_wasm_version: BuildVersion::default(),
-                new_wasm_version: result.canister_wasm.version,
-                new_wasm: WasmToInstall::Default(result.canister_wasm.module),
-                deposit_cycles_if_needed: true,
-                args: result.init_args,
-                mode: CanisterInstallMode::Install,
-                stop_start_canister: false,
-            })
-            .await
-            {
-                Ok(_) => {
-                    let response = mutate_state(|state| commit(args.canister_id, wasm_version, state));
-                    info!(canister_id = %args.canister_id, "local user index canister added");
-                    response
-                }
-                Err(error) => InternalError(format!("{error:?}")),
+
+            if let Err(error) = set_controllers(args.canister_id, vec![result.this_canister_id]).await {
+                InternalError(format!("Failed to set controller: {error:?}"))
+            } else if let Err(error) = install_basic(args.canister_id, result.canister_wasm, result.init_args).await {
+                InternalError(format!("Failed to install canister: {error:?}"))
+            } else {
+                let response = mutate_state(|state| commit(args.canister_id, wasm_version, state));
+                info!(canister_id = %args.canister_id, "local user index canister added");
+                response
             }
         }
         Err(response) => response,
     }
 }
 
 struct PrepareResult {
+    this_canister_id: CanisterId,
     canister_wasm: CanisterWasm,
     init_args: local_user_index_canister::init::Args,
 }
 
 fn prepare(args: &Args, state: &RuntimeState) -> Result<PrepareResult, Response> {
     if !state.data.local_index_map.contains_key(&args.canister_id) {
         Ok(PrepareResult {
+            this_canister_id: state.env.canister_id(),
             canister_wasm: state.data.local_user_index_canister_wasm_for_new_canisters.clone(),
             init_args: local_user_index_canister::init::Args {
                 user_canister_wasm: state.data.user_canister_wasm.clone(),

diff --git a/backend/libraries/utils/src/canister/create.rs b/backend/libraries/utils/src/canister/create.rs
@@ -1,10 +1,10 @@
-use crate::canister::{install, CanisterToInstall, WasmToInstall};
+use crate::canister::install_basic;
 use candid::{CandidType, Principal};
 use ic_cdk::api::call::{CallResult, RejectionCode};
 use ic_cdk::api::management_canister;
-use ic_cdk::api::management_canister::main::{CanisterInstallMode, CanisterSettings, CreateCanisterArgument};
+use ic_cdk::api::management_canister::main::{CanisterSettings, CreateCanisterArgument};
 use tracing::error;
-use types::{BuildVersion, CanisterId, CanisterWasm, Cycles};
+use types::{CanisterId, CanisterWasm, Cycles};
 
 #[derive(Debug)]
 pub enum CreateAndInstallError {
@@ -32,18 +32,7 @@ pub async fn create_and_install<A: CandidType>(
         },
     };
 
-    match install(CanisterToInstall {
-        canister_id,
-        current_wasm_version: BuildVersion::default(),
-        new_wasm_version: wasm.version,
-        new_wasm: WasmToInstall::Default(wasm.module),
-        deposit_cycles_if_needed: true,
-        args: init_args,
-        mode: CanisterInstallMode::Install,
-        stop_start_canister: false,
-    })
-    .await
-    {
+    match install_basic(canister_id, wasm, init_args).await {
         Ok(_) => Ok(canister_id),
         Err((code, msg)) => Err(CreateAndInstallError::InstallFailed(canister_id, code, msg)),
     }

diff --git a/backend/libraries/utils/src/canister/install.rs b/backend/libraries/utils/src/canister/install.rs
@@ -5,7 +5,7 @@ use ic_cdk::api::call::{CallResult, RejectionCode};
 use ic_cdk::api::management_canister;
 use ic_cdk::api::management_canister::main::{CanisterInstallMode, ChunkHash, InstallChunkedCodeArgument, InstallCodeArgument};
 use tracing::{error, trace};
-use types::{BuildVersion, CanisterId, Cycles, Hash};
+use types::{BuildVersion, CanisterId, CanisterWasm, Cycles, Hash};
 
 pub struct CanisterToInstall<A: CandidType> {
     pub canister_id: CanisterId,
@@ -29,6 +29,21 @@ pub struct ChunkedWasmToInstall {
     pub store_canister_id: CanisterId,
 }
 
+pub async fn install_basic<A: CandidType>(canister_id: CanisterId, wasm: CanisterWasm, init_args: A) -> CallResult<()> {
+    install(CanisterToInstall {
+        canister_id,
+        current_wasm_version: BuildVersion::default(),
+        new_wasm_version: wasm.version,
+        new_wasm: WasmToInstall::Default(wasm.module),
+        deposit_cycles_if_needed: true,
+        args: init_args,
+        mode: CanisterInstallMode::Install,
+        stop_start_canister: false,
+    })
+    .await
+    .map(|_| ())
+}
+
 pub async fn install<A: CandidType>(canister_to_install: CanisterToInstall<A>) -> CallResult<Option<Cycles>> {
     let canister_id = canister_to_install.canister_id;
     let mode = canister_to_install.mode;