From ab83fd7e0a31dd11605077e8d9b23d41a084f0dd Mon Sep 17 00:00:00 2001
From: Nikolay Komarevskiy
 <90605504+nikolay-komarevskiy@users.noreply.github.com>
Date: Mon, 4 Nov 2024 23:02:00 +0100
Subject: [PATCH] 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 <idx@dfinity.org>
---
 .../rate_limits/canister/add_config.rs        |  6 +++-
 .../rate_limits/canister/canister.rs          | 22 ++++++++++++---
 .../rate_limits/canister/state.rs             | 28 +++++++++----------
 3 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/rs/boundary_node/rate_limits/canister/add_config.rs b/rs/boundary_node/rate_limits/canister/add_config.rs
index f936541ec23..436db27797b 100644
--- a/rs/boundary_node/rate_limits/canister/add_config.rs
+++ b/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>;
diff --git a/rs/boundary_node/rate_limits/canister/canister.rs b/rs/boundary_node/rate_limits/canister/canister.rs
index ce7e8be08bb..22e1264b1ac 100644
--- a/rs/boundary_node/rate_limits/canister/canister.rs
+++ b/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]
diff --git a/rs/boundary_node/rate_limits/canister/state.rs b/rs/boundary_node/rate_limits/canister/state.rs
index cb6d75dbe42..c96c720ceb5 100644
--- a/rs/boundary_node/rate_limits/canister/state.rs
+++ b/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 {