feat(boundary): add inspect_message hook to canister

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, post_upgrade, query, update};
+use ic_cdk_macros::{init, inspect_message, post_upgrade, query, update};
 use rate_limits_api::{
     AddConfigResponse, ApiBoundaryNodeIdRecord, DiscloseRulesArg, DiscloseRulesResponse,
     GetApiBoundaryNodeIdsRequest, GetConfigResponse, GetRuleByIdResponse, InitArg, InputConfig,
@@ -22,6 +22,41 @@ use rate_limits_api::{
 const REGISTRY_CANISTER_ID: &str = "rwlgt-iiaaa-aaaaa-aaaaa-cai";
 const REGISTRY_CANISTER_METHOD: &str = "get_api_boundary_node_ids";
 
+const UPDATE_METHODS: [&str; 2] = ["add_config", "disclose_rules"];
+const REPLICATED_QUERY_METHOD: &str = "get_config";
+
+#[inspect_message]
+fn inspect_message() {
+    // In order for this hook to succeed, accept_message() must be invoked.
+    let caller_id: Principal = ic_cdk::api::caller();
+    let called_method = ic_cdk::api::call::method_name();
+
+    let (has_full_access, has_full_read_access) = with_canister_state(|state| {
+        let authorized_principal = state.get_authorized_principal();
+        (
+            Some(caller_id) == authorized_principal,
+            state.is_api_boundary_node_principal(&caller_id),
+        )
+    });
+
+    if called_method == REPLICATED_QUERY_METHOD {
+        if has_full_access || has_full_read_access {
+            ic_cdk::api::call::accept_message();
+        }
+    } else if UPDATE_METHODS.contains(&called_method.as_str()) {
+        if has_full_access {
+            ic_cdk::api::call::accept_message();
+        } else {
+            ic_cdk::api::trap("message_inspection_failed: unauthorized caller");
+        }
+    } else {
+        // All others calls are rejected
+        ic_cdk::api::trap(
+            "message_inspection_failed: method call is prohibited in the current context",
+        );
+    }
+}
+
 #[init]
 fn init(init_arg: InitArg) {
     ic_cdk::println!("Starting canister init");

rs/boundary_node/rate_limits/canister/state.rs

@@ -1,12 +1,13 @@
 use candid::Principal;
 use mockall::automock;
+use std::collections::HashSet;
 
 use crate::{
     add_config::{INIT_SCHEMA_VERSION, INIT_VERSION},
     storage::{
         LocalRef, StableMap, StorableConfig, StorableIncidentId, StorableIncidentMetadata,
         StorablePrincipal, StorableRuleId, StorableRuleMetadata, StorableVersion,
-        AUTHORIZED_PRINCIPAL, CONFIGS, INCIDENTS, RULES,
+        API_BOUNDARY_NODE_PRINCIPALS, AUTHORIZED_PRINCIPAL, CONFIGS, INCIDENTS, RULES,
     },
     types::{IncidentId, InputConfig, InputRule, RuleId, Timestamp, Version},
 };
@@ -31,6 +32,7 @@ pub trait CanisterApi {
         incident_id: IncidentId,
         rule_ids: StorableIncidentMetadata,
     ) -> Option<StorableIncidentMetadata>;
+    fn is_api_boundary_node_principal(&self, principal: &Principal) -> bool;
 }
 
 #[derive(Clone)]
@@ -39,6 +41,7 @@ pub struct CanisterState {
     rules: LocalRef<StableMap<StorableRuleId, StorableRuleMetadata>>,
     incidents: LocalRef<StableMap<StorableIncidentId, StorableIncidentMetadata>>,
     authorized_principal: LocalRef<StableMap<(), StorablePrincipal>>,
+    api_boundary_node_principals: LocalRef<HashSet<Principal>>,
 }
 
 impl CanisterState {
@@ -48,6 +51,7 @@ impl CanisterState {
             rules: &RULES,
             incidents: &INCIDENTS,
             authorized_principal: &AUTHORIZED_PRINCIPAL,
+            api_boundary_node_principals: &API_BOUNDARY_NODE_PRINCIPALS,
         }
     }
 }
@@ -128,6 +132,11 @@ impl CanisterApi for CanisterState {
                 .insert(StorableIncidentId(incident_id.0), rule_ids)
         })
     }
+
+    fn is_api_boundary_node_principal(&self, principal: &Principal) -> bool {
+        self.api_boundary_node_principals
+            .with(|cell| cell.borrow().contains(principal))
+    }
 }
 
 pub fn init_version_and_config(time: Timestamp, canister_api: impl CanisterApi) {

rs/boundary_node/rate_limits/canister_client/src/main.rs

@@ -236,9 +236,9 @@ async fn read_config(agent: &Agent, version: Version, canister_id: Principal) ->
     let args = Encode!(&Some(version)).unwrap();
 
     let response = agent
-        .update(&canister_id, "get_config")
+        .query(&canister_id, "get_config")
         .with_arg(args)
-        .call_and_wait()
+        .call()
         .await
         .expect("update call failed");
 
@@ -276,9 +276,9 @@ async fn read_rule(agent: &Agent, rule_id: &RuleId, canister_id: Principal) {
     let args = Encode!(rule_id).unwrap();
 
     let response = agent
-        .update(&canister_id, "get_rule_by_id")
+        .query(&canister_id, "get_rule_by_id")
         .with_arg(args)
-        .call_and_wait()
+        .call()
         .await
         .expect("update call failed");
 

rs/boundary_node/rate_limits/integration_tests/src/pocket_ic_helpers.rs

@@ -96,18 +96,26 @@ pub async fn get_installed_wasm_hash(pocket_ic: &PocketIc, canister_id: Principa
 pub async fn canister_call<R: DeserializeOwned + CandidType>(
     pocket_ic: &PocketIc,
     method: &str,
+    method_type: &str,
     canister_id: Principal,
     sender: Principal,
     payload: Vec<u8>,
 ) -> Result<R, String> {
-    let result = pocket_ic
-        .update_call(canister_id, sender, method, payload)
-        .await
-        .map_err(|err| err.to_string())?;
+    let result = match method_type {
+        "query" => pocket_ic
+            .query_call(canister_id, sender, method, payload)
+            .await
+            .map_err(|err| err.to_string())?,
+        "update" => pocket_ic
+            .update_call(canister_id, sender, method, payload)
+            .await
+            .map_err(|err| err.to_string())?,
+        _ => panic!("{method_type} is not allowed"),
+    };
 
     let result = match result {
         WasmResult::Reply(result) => result,
-        WasmResult::Reject(s) => panic!("Call to add_config failed: {:#?}", s),
+        WasmResult::Reject(s) => panic!("Call to {method} failed: {:#?}", s),
     };
 
     let decoded: R = Decode!(&result, R).unwrap();

rs/boundary_node/rate_limits/integration_tests/tests/rate_limit_canister_tests.rs

@@ -31,6 +31,7 @@ async fn main() {
     let response: GetConfigResponse = canister_call(
         &pocket_ic,
         "get_config",
+        "query",
         canister_id,
         Principal::anonymous(),
         input_version,
@@ -52,19 +53,22 @@ async fn main() {
     let response: AddConfigResponse = canister_call(
         &pocket_ic,
         "add_config",
+        "update",
         canister_id,
         Principal::anonymous(),
         input_config.clone(),
     )
-    .await
-    .unwrap();
+    .await;
+
+    let err_msg = response.unwrap_err();
 
-    assert!(response.unwrap_err().contains("Unauthorized"));
+    assert!(err_msg.contains("message_inspection_failed: unauthorized caller"));
 
     // Try add config using authorized principal as sender, assert success
     let response: AddConfigResponse = canister_call(
         &pocket_ic,
         "add_config",
+        "update",
         canister_id,
         authorized_principal,
         input_config,
@@ -80,6 +84,7 @@ async fn main() {
     let response: GetConfigResponse = canister_call(
         &pocket_ic,
         "get_config",
+        "query",
         canister_id,
         Principal::anonymous(),
         input_version,

rs/tests/boundary_nodes/rate_limit_canister_test.rs

@@ -325,9 +325,9 @@ async fn read_config(
     let args = Encode!(&Some(version)).unwrap();
 
     let response = agent
-        .update(&canister_id, "get_config")
+        .query(&canister_id, "get_config")
         .with_arg(args)
-        .call_and_wait()
+        .call()
         .await
         .expect("update call failed");
 
@@ -375,9 +375,9 @@ async fn read_rule(
     let args = Encode!(rule_id).unwrap();
 
     let response = agent
-        .update(&canister_id, "get_rule_by_id")
+        .query(&canister_id, "get_rule_by_id")
         .with_arg(args)
-        .call_and_wait()
+        .call()
         .await
         .expect("update call failed");