Skip to content

Commit

Permalink
chore: disallow replicated queries except get_config
Browse files Browse the repository at this point in the history
  • Loading branch information
IDX GitLab Automation committed Nov 11, 2024
1 parent dd1c4a2 commit b702ef3
Show file tree
Hide file tree
Showing 6 changed files with 68 additions and 35 deletions.
48 changes: 28 additions & 20 deletions rs/boundary_node/rate_limits/canister/canister.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,34 +22,42 @@ 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 CANISTER_UPDATE_METHODS: [&str; 2] = ["add_config", "disclose_rules"];
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();

// If the called method is not an update method, accept the message.
if !CANISTER_UPDATE_METHODS.contains(&called_method.as_str()) {
ic_cdk::api::call::accept_message();
} else {
// For the update methods:
// - Check if the canister's authorized principal is set
// - Check caller_id matches the authorized principal
with_canister_state(|state| {
if let Some(authorized_principal) = state.get_authorized_principal() {
if caller_id == authorized_principal {
ic_cdk::api::call::accept_message();
} else {
ic_cdk::api::trap("inspect_message_failed: unauthorized caller");
}
} else {
ic_cdk::api::trap(
"inspect_message_failed: authorized principal for canister is not set",
);
}
let (has_full_access, has_full_read_access) = with_canister_state(|state| {
let authorized_principal = state.get_authorized_principal().unwrap_or_else(|| {
ic_cdk::api::trap(
"message_inspection_failed: authorized principal for canister is not set",
)
});
(
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",
);
}
}

Expand Down
11 changes: 10 additions & 1 deletion rs/boundary_node/rate_limits/canister/state.rs
Original file line number Diff line number Diff line change
@@ -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},
};
Expand All @@ -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)]
Expand All @@ -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 {
Expand All @@ -48,6 +51,7 @@ impl CanisterState {
rules: &RULES,
incidents: &INCIDENTS,
authorized_principal: &AUTHORIZED_PRINCIPAL,
api_boundary_node_principals: &API_BOUNDARY_NODE_PRINCIPALS,
}
}
}
Expand Down Expand Up @@ -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) {
Expand Down
8 changes: 4 additions & 4 deletions rs/boundary_node/rate_limits/canister_client/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");

Expand Down Expand Up @@ -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");

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,18 +96,30 @@ 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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ async fn main() {
let response: GetConfigResponse = canister_call(
&pocket_ic,
"get_config",
"query",
canister_id,
Principal::anonymous(),
input_version,
Expand All @@ -52,6 +53,7 @@ async fn main() {
let response: AddConfigResponse = canister_call(
&pocket_ic,
"add_config",
"update",
canister_id,
Principal::anonymous(),
input_config.clone(),
Expand All @@ -60,12 +62,13 @@ async fn main() {

let err_msg = response.unwrap_err();

assert!(err_msg.contains("inspect_message_failed: unauthorized caller"));
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,
Expand All @@ -81,6 +84,7 @@ async fn main() {
let response: GetConfigResponse = canister_call(
&pocket_ic,
"get_config",
"query",
canister_id,
Principal::anonymous(),
input_version,
Expand Down
8 changes: 4 additions & 4 deletions rs/tests/boundary_nodes/rate_limit_canister_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");

Expand Down Expand Up @@ -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");

Expand Down

0 comments on commit b702ef3

Please sign in to comment.