From 8c04970d67df0963710b246b9b4ab8fe2443aba7 Mon Sep 17 00:00:00 2001 From: 0xZensh Date: Sat, 3 Aug 2024 10:23:28 +0800 Subject: [PATCH] chore: improve `idempotent-proxy-canister` caller states --- Cargo.lock | 6 +- Cargo.toml | 2 +- .../idempotent-proxy-canister.did | 19 +++--- src/idempotent-proxy-canister/src/api.rs | 67 ++++++++++++++----- .../src/api_admin.rs | 16 +++-- src/idempotent-proxy-canister/src/lib.rs | 4 +- src/idempotent-proxy-canister/src/store.rs | 30 +++++++-- 7 files changed, 104 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c2692ae..473f784 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1244,7 +1244,7 @@ dependencies = [ [[package]] name = "idempotent-proxy-canister" -version = "1.1.6" +version = "1.2.0" dependencies = [ "base64 0.22.1", "bytes", @@ -1265,7 +1265,7 @@ dependencies = [ [[package]] name = "idempotent-proxy-server" -version = "1.1.6" +version = "1.2.0" dependencies = [ "anyhow", "async-trait", @@ -1294,7 +1294,7 @@ dependencies = [ [[package]] name = "idempotent-proxy-types" -version = "1.1.6" +version = "1.2.0" dependencies = [ "base64 0.22.1", "ciborium", diff --git a/Cargo.toml b/Cargo.toml index 9230879..33dc751 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ strip = true opt-level = 's' [workspace.package] -version = "1.1.6" +version = "1.2.0" edition = "2021" repository = "https://github.com/ldclabs/idempotent-proxy" keywords = ["idempotent", "reverse", "proxy", "icp"] diff --git a/src/idempotent-proxy-canister/idempotent-proxy-canister.did b/src/idempotent-proxy-canister/idempotent-proxy-canister.did index 894c91b..f39fcb5 100644 --- a/src/idempotent-proxy-canister/idempotent-proxy-canister.did +++ b/src/idempotent-proxy-canister/idempotent-proxy-canister.did @@ -30,17 +30,18 @@ type InitArgs = record { }; type Result = variant { Ok : bool; Err : text }; type Result_1 = variant { Ok; Err : text }; -type Result_2 = variant { Ok : StateInfo; Err : text }; type StateInfo = record { - freezing_threshold : nat64; + proxy_token_public_key : text; + service_fee : nat64; ecdsa_key_name : text; managers : vec principal; - name : text; - auditors : vec principal; - schnorr_key_name : text; + cose : opt CoseClient; + uncollectible_cycles : nat; + callers : nat64; + agents : vec Agent; + incoming_cycles : nat; + proxy_token_refresh_interval : nat64; subnet_size : nat64; - namespace_total : nat64; - vetkd_key_name : text; }; type TransformArgs = record { context : blob; response : HttpResponse }; type TransformContext = record { @@ -60,13 +61,13 @@ service : (opt ChainArgs) -> { admin_remove_callers : (vec principal) -> (Result_1); admin_remove_managers : (vec principal) -> (Result_1); admin_set_agents : (vec Agent) -> (Result_1); - get_state : () -> (Result_2) query; - is_caller : (principal) -> (Result) query; + caller_info : (principal) -> (opt record { nat; nat64 }) query; parallel_call_all_ok : (CanisterHttpRequestArgument) -> (HttpResponse); parallel_call_any_ok : (CanisterHttpRequestArgument) -> (HttpResponse); parallel_call_cost : (CanisterHttpRequestArgument) -> (nat) query; proxy_http_request : (CanisterHttpRequestArgument) -> (HttpResponse); proxy_http_request_cost : (CanisterHttpRequestArgument) -> (nat) query; + state_info : () -> (StateInfo) query; validate_admin_add_managers : (vec principal) -> (Result_1); validate_admin_remove_managers : (vec principal) -> (Result_1); validate_admin_set_agents : (vec Agent) -> (Result_1); diff --git a/src/idempotent-proxy-canister/src/api.rs b/src/idempotent-proxy-canister/src/api.rs index 34cb80b..a883235 100644 --- a/src/idempotent-proxy-canister/src/api.rs +++ b/src/idempotent-proxy-canister/src/api.rs @@ -7,6 +7,8 @@ use std::collections::BTreeSet; use crate::{agent::Agent, cose::CoseClient, store}; +const MILLISECONDS: u64 = 1_000_000; + #[derive(CandidType, Deserialize, Serialize)] pub struct StateInfo { pub ecdsa_key_name: String, @@ -14,6 +16,7 @@ pub struct StateInfo { pub proxy_token_refresh_interval: u64, // seconds pub agents: Vec, pub managers: BTreeSet, + pub callers: u64, pub subnet_size: u64, pub service_fee: u64, // in cycles pub incoming_cycles: u128, @@ -22,8 +25,8 @@ pub struct StateInfo { } #[ic_cdk::query] -fn get_state() -> Result { - let s = store::state::with(|s| StateInfo { +fn state_info() -> StateInfo { + store::state::with(|s| StateInfo { ecdsa_key_name: s.ecdsa_key_name.clone(), proxy_token_public_key: s.proxy_token_public_key.clone(), proxy_token_refresh_interval: s.proxy_token_refresh_interval, @@ -38,18 +41,18 @@ fn get_state() -> Result { }) .collect(), managers: s.managers.clone(), + callers: s.callers.len() as u64, subnet_size: s.subnet_size, service_fee: s.service_fee, incoming_cycles: s.incoming_cycles, uncollectible_cycles: s.uncollectible_cycles, cose: s.cose.clone(), - }); - Ok(s) + }) } #[ic_cdk::query] -fn is_caller(id: Principal) -> Result { - store::state::with(|s| Ok(s.allowed_callers.contains(&id))) +fn caller_info(id: Principal) -> Option<(u128, u64)> { + store::state::with(|s| s.callers.get(&id).copied()) } #[ic_cdk::query] @@ -75,7 +78,8 @@ async fn parallel_call_cost(req: CanisterHttpRequestArgument) -> u128 { /// Proxy HTTP request by all agents in sequence until one returns an status <= 500 result. #[ic_cdk::update] async fn proxy_http_request(req: CanisterHttpRequestArgument) -> HttpResponse { - if !store::state::is_allowed(&ic_cdk::caller()) { + let caller = ic_cdk::caller(); + if !store::state::is_allowed(&caller) { return HttpResponse { status: Nat::from(403u64), body: "caller is not allowed".as_bytes().to_vec(), @@ -92,6 +96,7 @@ async fn proxy_http_request(req: CanisterHttpRequestArgument) -> HttpResponse { }; } + let balance = ic_cdk::api::call::msg_cycles_available128(); let calc = store::state::cycles_calculator(); store::state::receive_cycles( calc.ingress_cost(ic_cdk::api::call::arg_data_raw_size()), @@ -106,12 +111,22 @@ async fn proxy_http_request(req: CanisterHttpRequestArgument) -> HttpResponse { Ok(res) => { let cycles = calc.http_outcall_response_cost(calc.count_response_bytes(&res), 1); store::state::receive_cycles(cycles, true); + store::state::update_caller_state( + &caller, + balance - ic_cdk::api::call::msg_cycles_available128(), + ic_cdk::api::time() / MILLISECONDS, + ); return res; } Err(res) => last_err = Some(res), } } + store::state::update_caller_state( + &caller, + balance - ic_cdk::api::call::msg_cycles_available128(), + ic_cdk::api::time() / MILLISECONDS, + ); last_err.unwrap() } @@ -119,7 +134,8 @@ async fn proxy_http_request(req: CanisterHttpRequestArgument) -> HttpResponse { /// or a 500 HttpResponse with all result. #[ic_cdk::update] async fn parallel_call_all_ok(req: CanisterHttpRequestArgument) -> HttpResponse { - if !store::state::is_allowed(&ic_cdk::caller()) { + let caller = ic_cdk::caller(); + if !store::state::is_allowed(&caller) { return HttpResponse { status: Nat::from(403u64), body: "caller is not allowed".as_bytes().to_vec(), @@ -136,6 +152,7 @@ async fn parallel_call_all_ok(req: CanisterHttpRequestArgument) -> HttpResponse }; } + let balance = ic_cdk::api::call::msg_cycles_available128(); let calc = store::state::cycles_calculator(); let cycles = calc.ingress_cost(ic_cdk::api::call::arg_data_raw_size()) + calc.http_outcall_request_cost(calc.count_request_bytes(&req), agents.len()); @@ -143,7 +160,7 @@ async fn parallel_call_all_ok(req: CanisterHttpRequestArgument) -> HttpResponse let results = futures::future::try_join_all(agents.iter().map(|agent| agent.call(req.clone()))).await; - match results { + let result = match results { Err(res) => res, Ok(res) => { let mut results = res.into_iter(); @@ -164,22 +181,30 @@ async fn parallel_call_all_ok(req: CanisterHttpRequestArgument) -> HttpResponse let mut buf = vec![]; into_writer(&inconsistent_results, &mut buf) .expect("failed to encode inconsistent results"); - return HttpResponse { + HttpResponse { status: Nat::from(500u64), body: buf, headers: vec![], - }; + } + } else { + base_result } - - base_result } - } + }; + + store::state::update_caller_state( + &caller, + balance - ic_cdk::api::call::msg_cycles_available128(), + ic_cdk::api::time() / MILLISECONDS, + ); + result } /// Proxy HTTP request by all agents in parallel and return the first (status <= 500) result. #[ic_cdk::update] async fn parallel_call_any_ok(req: CanisterHttpRequestArgument) -> HttpResponse { - if !store::state::is_allowed(&ic_cdk::caller()) { + let caller = ic_cdk::caller(); + if !store::state::is_allowed(&caller) { return HttpResponse { status: Nat::from(403u64), body: "caller is not allowed".as_bytes().to_vec(), @@ -196,6 +221,7 @@ async fn parallel_call_any_ok(req: CanisterHttpRequestArgument) -> HttpResponse }; } + let balance = ic_cdk::api::call::msg_cycles_available128(); let calc = store::state::cycles_calculator(); let cycles = calc.ingress_cost(ic_cdk::api::call::arg_data_raw_size()) + calc.http_outcall_request_cost(calc.count_request_bytes(&req), agents.len()); @@ -204,7 +230,7 @@ async fn parallel_call_any_ok(req: CanisterHttpRequestArgument) -> HttpResponse let result = futures::future::select_ok(agents.iter().map(|agent| agent.call(req.clone()).boxed())) .await; - match result { + let result = match result { Ok((res, _)) => { let cycles = calc.http_outcall_response_cost(calc.count_response_bytes(&res), agents.len()); @@ -212,5 +238,12 @@ async fn parallel_call_any_ok(req: CanisterHttpRequestArgument) -> HttpResponse res } Err(res) => res, - } + }; + + store::state::update_caller_state( + &caller, + balance - ic_cdk::api::call::msg_cycles_available128(), + ic_cdk::api::time() / MILLISECONDS, + ); + result } diff --git a/src/idempotent-proxy-canister/src/api_admin.rs b/src/idempotent-proxy-canister/src/api_admin.rs index bc6a472..9d00c53 100644 --- a/src/idempotent-proxy-canister/src/api_admin.rs +++ b/src/idempotent-proxy-canister/src/api_admin.rs @@ -28,14 +28,20 @@ fn admin_add_caller(id: Principal) -> Result { Err("anonymous caller cannot be added".to_string())?; } - store::state::with_mut(|r| Ok(r.allowed_callers.insert(id))) + store::state::with_mut(|r| { + let has = r.callers.contains_key(&id); + r.callers.entry(id).or_insert((0, 0)); + Ok(!has) + }) } #[ic_cdk::update(guard = "is_controller_or_manager")] -fn admin_add_callers(mut args: BTreeSet) -> Result<(), String> { +fn admin_add_callers(args: BTreeSet) -> Result<(), String> { validate_principals(&args)?; store::state::with_mut(|r| { - r.allowed_callers.append(&mut args); + args.into_iter().for_each(|p| { + r.callers.entry(p).or_insert((0, 0)); + }); Ok(()) }) } @@ -44,7 +50,9 @@ fn admin_add_callers(mut args: BTreeSet) -> Result<(), String> { fn admin_remove_callers(args: BTreeSet) -> Result<(), String> { validate_principals(&args)?; store::state::with_mut(|r| { - r.allowed_callers.retain(|p| !args.contains(p)); + args.iter().for_each(|p| { + r.callers.remove(p); + }); Ok(()) }) } diff --git a/src/idempotent-proxy-canister/src/lib.rs b/src/idempotent-proxy-canister/src/lib.rs index 5835aac..1cfa294 100644 --- a/src/idempotent-proxy-canister/src/lib.rs +++ b/src/idempotent-proxy-canister/src/lib.rs @@ -1,6 +1,5 @@ use candid::Principal; use ic_cdk::api::management_canister::http_request::{CanisterHttpRequestArgument, HttpResponse}; -use ic_cose_types::types::state::StateInfo; use std::collections::BTreeSet; mod agent; @@ -13,7 +12,8 @@ mod init; mod store; mod tasks; -use crate::init::ChainArgs; +use api::StateInfo; +use init::ChainArgs; fn is_controller() -> Result<(), String> { let caller = ic_cdk::caller(); diff --git a/src/idempotent-proxy-canister/src/store.rs b/src/idempotent-proxy-canister/src/store.rs index 7adc1e9..efbee64 100644 --- a/src/idempotent-proxy-canister/src/store.rs +++ b/src/idempotent-proxy-canister/src/store.rs @@ -9,7 +9,11 @@ use ic_stable_structures::{ }; use serde::{Deserialize, Serialize}; use serde_bytes::ByteBuf; -use std::{borrow::Cow, cell::RefCell, collections::BTreeSet}; +use std::{ + borrow::Cow, + cell::RefCell, + collections::{BTreeMap, BTreeSet}, +}; use crate::{ agent::Agent, @@ -27,7 +31,9 @@ pub struct State { pub proxy_token_refresh_interval: u64, // seconds pub agents: Vec, pub managers: BTreeSet, - pub allowed_callers: BTreeSet, + pub allowed_callers: BTreeSet, //deprecated + #[serde(default)] + pub callers: BTreeMap, #[serde(default)] pub subnet_size: u64, #[serde(default)] @@ -154,7 +160,16 @@ pub mod state { } pub fn is_allowed(caller: &Principal) -> bool { - STATE.with(|r| r.borrow().allowed_callers.contains(caller)) + STATE.with(|r| r.borrow().callers.contains_key(caller)) + } + + pub fn update_caller_state(caller: &Principal, cycles: u128, now_ms: u64) { + STATE.with(|r| { + r.borrow_mut().callers.get_mut(caller).map(|v| { + v.0 = v.0.saturating_add(cycles); + v.1 = now_ms; + }) + }); } pub fn with(f: impl FnOnce(&State) -> R) -> R { @@ -185,7 +200,14 @@ pub mod state { pub fn load() { STATE_STORE.with(|r| { - let s = r.borrow_mut().get().clone(); + let mut s = r.borrow().get().clone(); + if !s.allowed_callers.is_empty() { + s.allowed_callers.iter().for_each(|p| { + s.callers.entry(*p).or_insert((0, 0)); + }); + s.allowed_callers.clear(); + } + STATE.with(|h| { *h.borrow_mut() = s; });