feat: support ic-cose

src/idempotent-proxy-canister/Cargo.toml

@@ -16,11 +16,6 @@ license.workspace = true
 crate-type = ["cdylib"]
 
 [dependencies]
-bytes = "1.6"
-candid = "0.10"
-ic-cdk = "0.14"
-ic-cdk-timers = "0.8"
-ic-stable-structures = "0.6"
 http = { workspace = true }
 base64 = { workspace = true }
 ciborium = { workspace = true }
@@ -29,3 +24,10 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 serde_bytes = { workspace = true }
 sha3 = { workspace = true }
+bytes = "1.6"
+candid = "0.10"
+ic-cdk = "0.14"
+ic-cdk-timers = "0.8"
+ic-stable-structures = "0.6"
+ic_cose_types = "0.1"
+getrandom = { version = "0.2", features = ["custom"] }

src/idempotent-proxy-canister/idempotent-proxy-canister.did

@@ -13,6 +13,7 @@ type CanisterHttpRequestArgument = record {
   headers : vec HttpHeader;
 };
 type ChainArgs = variant { Upgrade : UpgradeArgs; Init : InitArgs };
+type CoseClient = record { id : principal; namespace : text };
 type HttpHeader = record { value : text; name : text };
 type HttpMethod = variant { get; head; post };
 type HttpResponse = record {
@@ -23,18 +24,19 @@ type HttpResponse = record {
 type InitArgs = record {
   service_fee : nat64;
   ecdsa_key_name : text;
+  cose : opt CoseClient;
   proxy_token_refresh_interval : nat64;
   subnet_size : nat64;
 };
 type Result = variant { Ok : bool; Err : text };
 type Result_1 = variant { Ok; Err : text };
-type Result_2 = variant { Ok : State; Err };
-type State = record {
+type Result_2 = variant { Ok : StateInfo; Err };
+type StateInfo = record {
   proxy_token_public_key : text;
   service_fee : nat64;
   ecdsa_key_name : text;
   managers : vec principal;
-  allowed_callers : vec principal;
+  cose : opt CoseClient;
   uncollectible_cycles : nat;
   agents : vec Agent;
   incoming_cycles : nat;
@@ -48,6 +50,7 @@ type TransformContext = record {
 };
 type UpgradeArgs = record {
   service_fee : opt nat64;
+  cose : opt CoseClient;
   proxy_token_refresh_interval : opt nat64;
   subnet_size : opt nat64;
 };

src/idempotent-proxy-canister/src/cose.rs

@@ -0,0 +1,62 @@
+use candid::{utils::ArgumentEncoder, CandidType, Principal};
+use ic_cose_types::types::{PublicKeyInput, PublicKeyOutput, SignInput};
+use serde::{Deserialize, Serialize};
+use serde_bytes::ByteBuf;
+
+#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
+pub struct CoseClient {
+    pub id: Principal,
+    pub namespace: String,
+}
+
+impl CoseClient {
+    pub async fn ecdsa_public_key(&self, derivation_path: Vec<ByteBuf>) -> Result<ByteBuf, String> {
+        let output: Result<PublicKeyOutput, String> = call(
+            self.id,
+            "ecdsa_public_key",
+            (Some(PublicKeyInput {
+                ns: self.namespace.clone(),
+                derivation_path,
+            }),),
+            0,
+        )
+        .await?;
+        let output = output?;
+        Ok(output.public_key)
+    }
+
+    pub async fn ecdsa_sign(
+        &self,
+        derivation_path: Vec<ByteBuf>,
+        message: ByteBuf,
+    ) -> Result<ByteBuf, String> {
+        let output: Result<ByteBuf, String> = call(
+            self.id,
+            "ecdsa_sign",
+            (SignInput {
+                ns: self.namespace.clone(),
+                derivation_path,
+                message,
+            },),
+            0,
+        )
+        .await?;
+        output
+    }
+}
+
+async fn call<In, Out>(id: Principal, method: &str, args: In, cycles: u128) -> Result<Out, String>
+where
+    In: ArgumentEncoder + Send,
+    Out: candid::CandidType + for<'a> candid::Deserialize<'a>,
+{
+    let (res,): (Out,) = ic_cdk::api::call::call_with_payment128(id, method, args, cycles)
+        .await
+        .map_err(|(code, msg)| {
+            format!(
+                "failed to call {} on {:?}, code: {}, message: {}",
+                method, &id, code as u32, msg
+            )
+        })?;
+    Ok(res)
+}

src/idempotent-proxy-canister/src/ecdsa.rs

@@ -1,31 +1,4 @@
-use base64::{engine::general_purpose::URL_SAFE_NO_PAD as base64_url, Engine};
-use ciborium::into_writer;
 use ic_cdk::api::management_canister::ecdsa;
-use serde_bytes::ByteBuf;
-use sha3::{Digest, Sha3_256};
-
-// use Idempotent Proxy's Token: Token(pub u64, pub String, pub ByteBuf);
-// https://github.com/ldclabs/idempotent-proxy/blob/main/src/idempotent-proxy-types/src/auth.rs#L15
-pub async fn sign_proxy_token(
-    key_name: &str,
-    expire_at: u64, // UNIX timestamp, in seconds
-    message: &str,  // use RPCAgent.name as message
-) -> Result<String, String> {
-    let mut buf: Vec<u8> = Vec::new();
-    into_writer(&(expire_at, message), &mut buf).expect("failed to encode Token in CBOR format");
-    let digest = sha3_256(&buf);
-    let sig = sign_with(key_name, vec![b"sign_proxy_token".to_vec()], digest)
-        .await
-        .map_err(err_string)?;
-    buf.clear();
-    into_writer(&(expire_at, message, ByteBuf::from(sig)), &mut buf).map_err(err_string)?;
-    Ok(base64_url.encode(buf))
-}
-
-pub async fn get_proxy_token_public_key(key_name: &str) -> Result<String, String> {
-    let pk = public_key_with(key_name, vec![b"sign_proxy_token".to_vec()]).await?;
-    Ok(base64_url.encode(pk.public_key))
-}
 
 pub async fn sign_with(
     key_name: &str,
@@ -67,13 +40,3 @@ pub async fn public_key_with(
 
     Ok(response)
 }
-
-pub fn err_string(err: impl std::fmt::Display) -> String {
-    err.to_string()
-}
-
-pub fn sha3_256(data: &[u8]) -> [u8; 32] {
-    let mut hasher = Sha3_256::new();
-    hasher.update(data);
-    hasher.finalize().into()
-}

src/idempotent-proxy-canister/src/init.rs

@@ -2,7 +2,7 @@ use candid::CandidType;
 use serde::Deserialize;
 use std::time::Duration;
 
-use crate::{store, tasks};
+use crate::{cose::CoseClient, store, tasks};
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub enum ChainArgs {
@@ -16,13 +16,15 @@ pub struct InitArgs {
     proxy_token_refresh_interval: u64, // seconds
     subnet_size: u64,       // set to 0 to disable receiving cycles
     service_fee: u64,       // in cycles
+    cose: Option<CoseClient>,
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub struct UpgradeArgs {
     proxy_token_refresh_interval: Option<u64>, // seconds
     subnet_size: Option<u64>,
     service_fee: Option<u64>, // in cycles
+    cose: Option<CoseClient>,
 }
 
 #[ic_cdk::init]
@@ -42,6 +44,7 @@ fn init(args: Option<ChainArgs>) {
                 } else {
                     100_000_000
                 };
+                s.cose = args.cose;
             });
         }
         ChainArgs::Upgrade(_) => {
@@ -89,6 +92,9 @@ fn post_upgrade(args: Option<ChainArgs>) {
                 if let Some(service_fee) = args.service_fee {
                     s.service_fee = service_fee;
                 }
+                if let Some(cose) = args.cose {
+                    s.cose = Some(cose);
+                }
             });
         }
         Some(ChainArgs::Init(_)) => {
@@ -101,6 +107,7 @@ fn post_upgrade(args: Option<ChainArgs>) {
 
     ic_cdk_timers::set_timer(Duration::from_secs(0), || {
         ic_cdk::spawn(async {
+            store::state::init_ecdsa_public_key().await;
             tasks::refresh_proxy_token().await;
         })
     });

src/idempotent-proxy-canister/src/lib.rs

@@ -1,17 +1,19 @@
-use candid::{Nat, Principal};
+use candid::{CandidType, Nat, Principal};
 use ciborium::into_writer;
 use futures::FutureExt;
 use ic_cdk::api::management_canister::http_request::{CanisterHttpRequestArgument, HttpResponse};
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeSet;
 
 mod agent;
+mod cose;
 mod cycles;
 mod ecdsa;
 mod init;
 mod store;
 mod tasks;
 
-use crate::init::ChainArgs;
+use crate::{agent::Agent, cose::CoseClient, init::ChainArgs};
 
 static ANONYMOUS: Principal = Principal::anonymous();
 
@@ -39,9 +41,9 @@ fn validate_admin_set_managers(args: BTreeSet<Principal>) -> Result<(), String>
 async fn admin_set_agents(agents: Vec<agent::Agent>) -> Result<(), String> {
     validate_admin_set_agents(agents.clone())?;
 
-    let (ecdsa_key_name, proxy_token_refresh_interval) =
-        store::state::with(|s| (s.ecdsa_key_name.clone(), s.proxy_token_refresh_interval));
-    tasks::update_proxy_token(ecdsa_key_name, proxy_token_refresh_interval, agents).await;
+    let (signer, proxy_token_refresh_interval) =
+        store::state::with(|s| (s.signer(), s.proxy_token_refresh_interval));
+    tasks::update_proxy_token(signer, proxy_token_refresh_interval, agents).await;
     Ok(())
 }
 
@@ -64,14 +66,43 @@ fn admin_remove_caller(id: Principal) -> Result<bool, String> {
     store::state::with_mut(|r| Ok(r.allowed_callers.remove(&id)))
 }
 
+#[derive(CandidType, Deserialize, Serialize)]
+pub struct StateInfo {
+    pub ecdsa_key_name: String,
+    pub proxy_token_public_key: String,
+    pub proxy_token_refresh_interval: u64, // seconds
+    pub agents: Vec<Agent>,
+    pub managers: BTreeSet<Principal>,
+    pub subnet_size: u64,
+    pub service_fee: u64, // in cycles
+    pub incoming_cycles: u128,
+    pub uncollectible_cycles: u128,
+    pub cose: Option<CoseClient>,
+}
+
 #[ic_cdk::query]
-fn get_state() -> Result<store::State, ()> {
-    let mut s = store::state::with(|s| s.clone());
-    if is_controller_or_manager().is_err() {
-        s.agents.iter_mut().for_each(|a| {
-            a.proxy_token = None;
-        })
-    }
+fn get_state() -> Result<StateInfo, ()> {
+    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(),
+        proxy_token_refresh_interval: s.proxy_token_refresh_interval,
+        agents: s
+            .agents
+            .iter()
+            .map(|a| Agent {
+                name: a.name.clone(),
+                endpoint: a.endpoint.clone(),
+                max_cycles: a.max_cycles,
+                proxy_token: a.proxy_token.clone(),
+            })
+            .collect(),
+        managers: s.managers.clone(),
+        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)
 }
 
@@ -256,4 +287,21 @@ fn is_controller_or_manager() -> Result<(), String> {
     }
 }
 
+#[cfg(all(
+    target_arch = "wasm32",
+    target_vendor = "unknown",
+    target_os = "unknown"
+))]
+/// A getrandom implementation that always fails
+pub fn always_fail(_buf: &mut [u8]) -> Result<(), getrandom::Error> {
+    Err(getrandom::Error::UNSUPPORTED)
+}
+
+#[cfg(all(
+    target_arch = "wasm32",
+    target_vendor = "unknown",
+    target_os = "unknown"
+))]
+getrandom::register_custom_getrandom!(always_fail);
+
 ic_cdk::export_candid!();