feat(boundary): add rate-limit canister post_upgrade (#2374)

This `post_upgrade` hook serves the following purposes:
- Allows changing the authorized principal during canister upgrade phase
(upgrades will be done via nns proposals)
- Resets periodic timer task of polling API boundary nodes, otherwise
this task remains stopped after upgrade

Co-authored-by: IDX GitLab Automation <[email protected]>

rs/boundary_node/rate_limits/canister/add_config.rs

@@ -1,4 +1,7 @@
-use crate::{storage::StorableIncidentMetadata, types::Timestamp};
+use crate::{
+    storage::StorableIncidentMetadata,
+    types::{SchemaVersion, Timestamp},
+};
 use anyhow::Context;
 use getrandom::getrandom;
 use rate_limits_api::IncidentId;
@@ -15,6 +18,7 @@ use crate::{
 };
 
 pub const INIT_VERSION: Version = 1;
+pub const INIT_SCHEMA_VERSION: SchemaVersion = 1;
 
 pub trait AddsConfig {
     fn add_config(&self, config: InputConfig, time: Timestamp) -> Result<(), AddConfigError>;

rs/boundary_node/rate_limits/canister/canister.rs

@@ -12,7 +12,7 @@ use crate::storage::API_BOUNDARY_NODE_PRINCIPALS;
 use candid::Principal;
 use ic_canisters_http_types::{HttpRequest, HttpResponse, HttpResponseBuilder};
 use ic_cdk::api::call::call;
-use ic_cdk_macros::{init, query, update};
+use ic_cdk_macros::{init, post_upgrade, query, update};
 use rate_limits_api::{
     AddConfigResponse, ApiBoundaryNodeIdRecord, DiscloseRulesArg, DiscloseRulesResponse,
     GetApiBoundaryNodeIdsRequest, GetConfigResponse, GetRuleByIdResponse, InitArg, InputConfig,
@@ -22,7 +22,6 @@ use rate_limits_api::{
 const REGISTRY_CANISTER_ID: &str = "rwlgt-iiaaa-aaaaa-aaaaa-cai";
 const REGISTRY_CANISTER_METHOD: &str = "get_api_boundary_node_ids";
 
-// TODO: implement proper canister lifecycle: upgrade, post_upgrade, ...
 #[init]
 fn init(init_arg: InitArg) {
     ic_cdk::println!("Starting canister init");
@@ -32,11 +31,26 @@ fn init(init_arg: InitArg) {
             state.set_authorized_principal(principal);
         });
     }
-    // Initialize an empty config corresponding to version=1
-    init_version_and_config(1);
+    with_canister_state(|state| {
+        if state.get_version().is_none() {
+            ic_cdk::println!("Initializing rate-limit config");
+            let current_time = ic_cdk::api::time();
+            init_version_and_config(current_time, state);
+        } else {
+            ic_cdk::println!("Rate-limit config is already initialized");
+        }
+    });
     // Spawn periodic job of fetching latest API boundary node topology
     // API boundary nodes are authorized readers of all config rules (including not yet disclosed ones)
     periodically_poll_api_boundary_nodes(init_arg.registry_polling_period_secs);
+    ic_cdk::println!("Finished canister init");
+}
+
+#[post_upgrade]
+fn post_upgrade(init_arg: InitArg) {
+    ic_cdk::println!("Starting canister post-upgrade");
+    init(init_arg);
+    ic_cdk::println!("Finished canister post-upgrade");
 }
 
 #[query]

rs/boundary_node/rate_limits/canister/state.rs

@@ -1,16 +1,15 @@
 use candid::Principal;
-use ic_cdk::api::time;
 use mockall::automock;
 use rate_limits_api::IncidentId;
 
 use crate::{
-    add_config::INIT_VERSION,
+    add_config::{INIT_SCHEMA_VERSION, INIT_VERSION},
     storage::{
         LocalRef, StableMap, StorableConfig, StorableIncidentId, StorableIncidentMetadata,
         StorablePrincipal, StorableRuleId, StorableRuleMetadata, StorableVersion,
         AUTHORIZED_PRINCIPAL, CONFIGS, INCIDENTS, RULES,
     },
-    types::{InputConfig, InputRule, RuleId, Version},
+    types::{InputConfig, InputRule, RuleId, Timestamp, Version},
 };
 
 #[automock]
@@ -167,18 +166,17 @@ impl CanisterApi for CanisterState {
     }
 }
 
-pub fn init_version_and_config(version: Version) {
-    with_canister_state(|state| {
-        let config = StorableConfig {
-            schema_version: 1,
-            active_since: time(),
-            rule_ids: vec![],
-        };
-        assert!(
-            state.add_config(INIT_VERSION, config),
-            "Config for version={version} already exists!"
-        );
-    })
+pub fn init_version_and_config(time: Timestamp, canister_api: impl CanisterApi) {
+    // Initialize config with an empty vector of rules
+    let config = StorableConfig {
+        schema_version: INIT_SCHEMA_VERSION,
+        active_since: time,
+        rule_ids: vec![],
+    };
+    assert!(
+        canister_api.add_config(INIT_VERSION, config),
+        "Config for version={INIT_VERSION} already exists!"
+    );
 }
 
 pub fn with_canister_state<R>(f: impl FnOnce(CanisterState) -> R) -> R {