Automatically add local indexes to CyclesDispenser's allow list (#7070)

open-chat-labs · Dec 16, 2024 · 02e8b33 · 02e8b33
1 parent d3f9d43
commit 02e8b33
Show file tree

Hide file tree

Showing 20 changed files with 76 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/backend/canisters/cycles_dispenser/CHANGELOG.md b/backend/canisters/cycles_dispenser/CHANGELOG.md
@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Allow SNS controlled canisters to add additional canisters to the allow list ([#7070](https://github.com/open-chat-labs/open-chat/pull/7070))
+
 ## [[2.0.1511](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1511-cycles_dispenser)] - 2024-12-13
 
 ### Changed

diff --git a/backend/canisters/cycles_dispenser/impl/src/guards.rs b/backend/canisters/cycles_dispenser/impl/src/guards.rs
@@ -7,3 +7,11 @@ pub fn caller_is_governance_principal() -> Result<(), String> {
         Err("Caller is not a governance principal".to_string())
     }
 }
+
+pub fn caller_is_authorized_to_add_canister() -> Result<(), String> {
+    if read_state(|state| state.is_caller_authorized_to_add_canister()) {
+        Ok(())
+    } else {
+        Err("Caller is not authorized to add a canister".to_string())
+    }
+}
diff --git a/backend/canisters/cycles_dispenser/impl/src/jobs/top_up_sns_canisters.rs b/backend/canisters/cycles_dispenser/impl/src/jobs/top_up_sns_canisters.rs
@@ -36,7 +36,8 @@ async fn run_async(canister_id: CanisterId) {
         // Add SNS canisters to the whitelist
         mutate_state(|state| {
             let now = state.env.now();
-            for canister_id in canisters.iter().flat_map(|c| c.canister_id) {
+            for canister_id in canisters.iter().filter_map(|s| s.canister_id) {
+                state.data.canisters_directly_controlled_by_sns_root.insert(canister_id);
                 state.data.canisters.add(canister_id, now);
             }
         });

diff --git a/backend/canisters/cycles_dispenser/impl/src/lib.rs b/backend/canisters/cycles_dispenser/impl/src/lib.rs
@@ -5,7 +5,7 @@ use ic_ledger_types::{BlockIndex, Tokens};
 use ledger_utils::default_ledger_account;
 use serde::{Deserialize, Serialize};
 use std::cell::RefCell;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 use types::{BuildVersion, CanisterId, Cycles, Milliseconds, TimestampMillis, Timestamped};
 use utils::env::Environment;
 
@@ -37,6 +37,12 @@ impl State {
         self.data.governance_principals.contains(&self.env.caller())
     }
 
+    pub fn is_caller_authorized_to_add_canister(&self) -> bool {
+        let caller = self.env.caller();
+        self.data.governance_principals.contains(&caller)
+            || self.data.canisters_directly_controlled_by_sns_root.contains(&caller)
+    }
+
     pub fn metrics(&self) -> Metrics {
         Metrics {
             heap_memory_used: utils::memory::heap(),
@@ -64,6 +70,8 @@ impl State {
 struct Data {
     pub governance_principals: HashSet<Principal>,
     pub canisters: Canisters,
+    #[serde(default)]
+    pub canisters_directly_controlled_by_sns_root: BTreeSet<CanisterId>,
     pub sns_root_canister: Option<CanisterId>,
     pub max_top_up_amount: Cycles,
     pub min_interval: Milliseconds,
@@ -93,6 +101,7 @@ impl Data {
         Data {
             governance_principals: governance_principals.into_iter().collect(),
             canisters: Canisters::new(canisters, now),
+            canisters_directly_controlled_by_sns_root: BTreeSet::default(),
             sns_root_canister: None,
             max_top_up_amount,
             min_interval,

diff --git a/backend/canisters/cycles_dispenser/impl/src/updates/add_canister.rs b/backend/canisters/cycles_dispenser/impl/src/updates/add_canister.rs
@@ -1,10 +1,10 @@
-use crate::guards::caller_is_governance_principal;
+use crate::guards::caller_is_authorized_to_add_canister;
 use crate::{mutate_state, State};
 use canister_api_macros::proposal;
 use canister_tracing_macros::trace;
 use cycles_dispenser_canister::add_canister::{Response::*, *};
 
-#[proposal(guard = "caller_is_governance_principal")]
+#[proposal(guard = "caller_is_authorized_to_add_canister")]
 #[trace]
 fn add_canister(args: Args) -> Response {
     mutate_state(|state| add_canister_impl(args, state))

diff --git a/backend/canisters/group_index/CHANGELOG.md b/backend/canisters/group_index/CHANGELOG.md
@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Add new LocalGroupIndexes to the CyclesDispenser's allow list ([#7070](https://github.com/open-chat-labs/open-chat/pull/7070))
+
 ## [[2.0.1509](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1509-group_index)] - 2024-12-13
 
 ### Changed

diff --git a/backend/canisters/group_index/impl/Cargo.toml b/backend/canisters/group_index/impl/Cargo.toml
@@ -19,6 +19,7 @@ canister_tracing_macros = { path = "../../../libraries/canister_tracing_macros"
 community_canister = { path = "../../community/api" }
 community_canister_c2c_client = { path = "../../community/c2c_client" }
 constants = { path = "../../../libraries/constants" }
+cycles_dispenser_canister = { path = "../../cycles_dispenser/api" }
 fire_and_forget_handler = { path = "../../../libraries/fire_and_forget_handler" }
 futures = { workspace = true }
 group_canister = { path = "../../group/api" }

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
@@ -73,6 +73,12 @@ fn prepare(args: &Args, state: &RuntimeState) -> Result<PrepareResult, Response>
 
 fn commit(canister_id: CanisterId, wasm_version: BuildVersion, state: &mut RuntimeState) -> Response {
     if state.data.local_index_map.add_index(canister_id, wasm_version) {
+        state.data.fire_and_forget_handler.send_candid(
+            state.data.cycles_dispenser_canister_id,
+            "add_canister",
+            cycles_dispenser_canister::add_canister::Args { canister_id },
+        );
+
         Success
     } else {
         AlreadyAdded

diff --git a/backend/canisters/notifications_index/CHANGELOG.md b/backend/canisters/notifications_index/CHANGELOG.md
@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Add new Notifications canisters to the CyclesDispenser's allow list ([#7070](https://github.com/open-chat-labs/open-chat/pull/7070))
+
 ## [[2.0.1518](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1518-notifications_index)] - 2024-12-13
 
 ### Added

diff --git a/backend/canisters/notifications_index/impl/Cargo.toml b/backend/canisters/notifications_index/impl/Cargo.toml
@@ -16,6 +16,8 @@ canister_logger = { path = "../../../libraries/canister_logger" }
 canister_state_macros = { path = "../../../libraries/canister_state_macros" }
 canister_tracing_macros = { path = "../../../libraries/canister_tracing_macros" }
 constants = { path = "../../../libraries/constants" }
+cycles_dispenser_canister = { path = "../../cycles_dispenser/api" }
+fire_and_forget_handler = { path = "../../../libraries/fire_and_forget_handler" }
 futures = { workspace = true }
 http_request = { path = "../../../libraries/http_request" }
 human_readable = { path = "../../../libraries/human_readable" }

diff --git a/backend/canisters/notifications_index/impl/src/lib.rs b/backend/canisters/notifications_index/impl/src/lib.rs
@@ -2,6 +2,7 @@ use crate::model::notifications_canister::NotificationsCanister;
 use crate::model::subscriptions::Subscriptions;
 use candid::Principal;
 use canister_state_macros::canister_state;
+use fire_and_forget_handler::FireAndForgetHandler;
 use notifications_index_canister::{NotificationsIndexEvent, SubscriptionAdded, SubscriptionRemoved};
 use principal_to_user_id_map::PrincipalToUserIdMap;
 use serde::{Deserialize, Serialize};
@@ -125,6 +126,8 @@ struct Data {
     pub notifications_canister_wasm_for_upgrades: CanisterWasm,
     pub canisters_requiring_upgrade: CanistersRequiringUpgrade,
     pub notifications_index_event_sync_queue: CanisterEventSyncQueue<NotificationsIndexEvent>,
+    #[serde(default)]
+    pub fire_and_forget_handler: FireAndForgetHandler,
     pub rng_seed: [u8; 32],
     pub test_mode: bool,
 }
@@ -150,6 +153,7 @@ impl Data {
             notifications_canister_wasm_for_upgrades: notifications_canister_wasm,
             canisters_requiring_upgrade: CanistersRequiringUpgrade::default(),
             notifications_index_event_sync_queue: CanisterEventSyncQueue::default(),
+            fire_and_forget_handler: FireAndForgetHandler::default(),
             rng_seed: [0; 32],
             test_mode,
         }

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
@@ -69,6 +69,12 @@ fn commit(canister_id: CanisterId, wasm_version: BuildVersion, state: &mut Runti
             );
         }
 
+        state.data.fire_and_forget_handler.send_candid(
+            state.data.cycles_dispenser_canister_id,
+            "add_canister",
+            cycles_dispenser_canister::add_canister::Args { canister_id },
+        );
+
         Success
     } else {
         AlreadyAdded

diff --git a/backend/canisters/user_index/CHANGELOG.md b/backend/canisters/user_index/CHANGELOG.md
@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Add new LocalUserIndexes to the CyclesDispenser's allow list ([#7070](https://github.com/open-chat-labs/open-chat/pull/7070))
+
 ## [[2.0.1508](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1508-user_index)] - 2024-12-13
 
 ### Added

diff --git a/backend/canisters/user_index/impl/Cargo.toml b/backend/canisters/user_index/impl/Cargo.toml
@@ -20,13 +20,14 @@ chat_events = { path = "../../../libraries/chat_events" }
 community_canister = { path = "../../community/api" }
 community_canister_c2c_client = { path = "../../community/c2c_client" }
 constants = { path = "../../../libraries/constants" }
+cycles_dispenser_canister = { path = "../../cycles_dispenser/api" }
 dataurl = { workspace = true }
 event_store_producer = { workspace = true, features = ["json"] }
 event_store_producer_cdk_runtime = { workspace = true }
+fire_and_forget_handler = { path = "../../../libraries/fire_and_forget_handler" }
 futures = { workspace = true }
 group_canister = { path = "../../group/api" }
 group_canister_c2c_client = { path = "../../group/c2c_client" }
-fire_and_forget_handler = { path = "../../../libraries/fire_and_forget_handler" }
 hex = { workspace = true }
 http_request = { path = "../../../libraries/http_request" }
 human_readable = { path = "../../../libraries/human_readable" }

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
@@ -94,6 +94,12 @@ fn commit(canister_id: CanisterId, wasm_version: BuildVersion, state: &mut Runti
         }
         crate::jobs::sync_events_to_local_user_index_canisters::try_run_now(state);
 
+        state.data.fire_and_forget_handler.send_candid(
+            state.data.cycles_dispenser_canister_id,
+            "add_canister",
+            cycles_dispenser_canister::add_canister::Args { canister_id },
+        );
+
         Success
     } else {
         AlreadyAdded

diff --git a/backend/integration_tests/src/setup.rs b/backend/integration_tests/src/setup.rs
@@ -306,9 +306,6 @@ fn install_canisters(env: &mut PocketIc, controller: Principal) -> CanisterIds {
             user_index_canister_id,
             group_index_canister_id,
             notifications_index_canister_id,
-            local_user_index_canister_id,
-            local_group_index_canister_id,
-            notifications_canister_id,
             online_users_canister_id,
             proposals_bot_canister_id,
             storage_index_canister_id,

diff --git a/backend/libraries/fire_and_forget_handler/Cargo.toml b/backend/libraries/fire_and_forget_handler/Cargo.toml
@@ -6,6 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+candid = { workspace = true }
 canister_client = { path = "../canister_client" }
 canister_time = { path = "../canister_time" }
 constants = { path = "../constants" }

diff --git a/backend/libraries/fire_and_forget_handler/src/lib.rs b/backend/libraries/fire_and_forget_handler/src/lib.rs
@@ -1,3 +1,4 @@
+use candid::CandidType;
 use canister_client::make_c2c_call_raw;
 use constants::SECOND_IN_MS;
 use ic_cdk_timers::TimerId;
@@ -34,6 +35,10 @@ impl FireAndForgetHandler {
         ic_cdk::spawn(self.clone().process_single(call));
     }
 
+    pub fn send_candid<A: CandidType>(&self, canister_id: CanisterId, method_name: impl Into<String>, args: A) {
+        self.send(canister_id, method_name, candid::encode_one(args).unwrap());
+    }
+
     fn init(inner: FireAndForgetHandlerInner) -> Self {
         let wrapped = Rc::new(Mutex::new(inner));
         let handler = FireAndForgetHandler { inner: wrapped };

diff --git a/backend/tools/canister_installer/src/lib.rs b/backend/tools/canister_installer/src/lib.rs
@@ -195,9 +195,6 @@ async fn install_service_canisters_impl(
             canister_ids.user_index,
             canister_ids.group_index,
             canister_ids.notifications_index,
-            canister_ids.local_user_index,
-            canister_ids.local_group_index,
-            canister_ids.notifications,
             canister_ids.online_users,
             canister_ids.proposals_bot,
             canister_ids.storage_index,