From 3b64893abfadbcdab5dbd5f16f78b57cda9db2fb Mon Sep 17 00:00:00 2001
From: 0xZensh <txr1883@gmail.com>
Date: Sun, 28 Jul 2024 16:44:32 +0800
Subject: [PATCH] chore: improve API

---
 .../idempotent-proxy-canister.did             | 25 +++++++++++--------
 src/idempotent-proxy-canister/src/api.rs      |  7 +++++-
 .../src/api_admin.rs                          | 11 +++++++-
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/src/idempotent-proxy-canister/idempotent-proxy-canister.did b/src/idempotent-proxy-canister/idempotent-proxy-canister.did
index 187f7a9..894c91b 100644
--- a/src/idempotent-proxy-canister/idempotent-proxy-canister.did
+++ b/src/idempotent-proxy-canister/idempotent-proxy-canister.did
@@ -28,8 +28,9 @@ type InitArgs = record {
   proxy_token_refresh_interval : nat64;
   subnet_size : nat64;
 };
-type Result = variant { Ok; Err : text };
-type Result_1 = variant { Ok : StateInfo; Err };
+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;
   ecdsa_key_name : text;
@@ -53,18 +54,20 @@ type UpgradeArgs = record {
   subnet_size : opt nat64;
 };
 service : (opt ChainArgs) -> {
-  admin_add_callers : (vec principal) -> (Result);
-  admin_add_managers : (vec principal) -> (Result);
-  admin_remove_callers : (vec principal) -> (Result);
-  admin_remove_managers : (vec principal) -> (Result);
-  admin_set_agents : (vec Agent) -> (Result);
-  get_state : () -> (Result_1) query;
+  admin_add_caller : (principal) -> (Result);
+  admin_add_callers : (vec principal) -> (Result_1);
+  admin_add_managers : (vec principal) -> (Result_1);
+  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;
   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;
-  validate_admin_add_managers : (vec principal) -> (Result);
-  validate_admin_remove_managers : (vec principal) -> (Result);
-  validate_admin_set_agents : (vec Agent) -> (Result);
+  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 5fbb082..34cb80b 100644
--- a/src/idempotent-proxy-canister/src/api.rs
+++ b/src/idempotent-proxy-canister/src/api.rs
@@ -22,7 +22,7 @@ pub struct StateInfo {
 }
 
 #[ic_cdk::query]
-fn get_state() -> Result<StateInfo, ()> {
+fn get_state() -> Result<StateInfo, String> {
     let s = store::state::with(|s| StateInfo {
         ecdsa_key_name: s.ecdsa_key_name.clone(),
         proxy_token_public_key: s.proxy_token_public_key.clone(),
@@ -47,6 +47,11 @@ fn get_state() -> Result<StateInfo, ()> {
     Ok(s)
 }
 
+#[ic_cdk::query]
+fn is_caller(id: Principal) -> Result<bool, String> {
+    store::state::with(|s| Ok(s.allowed_callers.contains(&id)))
+}
+
 #[ic_cdk::query]
 async fn proxy_http_request_cost(req: CanisterHttpRequestArgument) -> u128 {
     let calc = store::state::cycles_calculator();
diff --git a/src/idempotent-proxy-canister/src/api_admin.rs b/src/idempotent-proxy-canister/src/api_admin.rs
index 5dfa593..bc6a472 100644
--- a/src/idempotent-proxy-canister/src/api_admin.rs
+++ b/src/idempotent-proxy-canister/src/api_admin.rs
@@ -1,5 +1,5 @@
 use candid::Principal;
-use ic_cose_types::validate_principals;
+use ic_cose_types::{validate_principals, ANONYMOUS};
 use std::collections::BTreeSet;
 
 use crate::{agent, is_controller, is_controller_or_manager, store, tasks};
@@ -22,6 +22,15 @@ fn admin_remove_managers(args: BTreeSet<Principal>) -> Result<(), String> {
     })
 }
 
+#[ic_cdk::update(guard = "is_controller_or_manager")]
+fn admin_add_caller(id: Principal) -> Result<bool, String> {
+    if id == ANONYMOUS {
+        Err("anonymous caller cannot be added".to_string())?;
+    }
+
+    store::state::with_mut(|r| Ok(r.allowed_callers.insert(id)))
+}
+
 #[ic_cdk::update(guard = "is_controller_or_manager")]
 fn admin_add_callers(mut args: BTreeSet<Principal>) -> Result<(), String> {
     validate_principals(&args)?;