diff --git a/rust/basic_bitcoin/src/basic_bitcoin/src/ecdsa_api.rs b/rust/basic_bitcoin/src/basic_bitcoin/src/ecdsa_api.rs index b50593137..615e46fb4 100644 --- a/rust/basic_bitcoin/src/basic_bitcoin/src/ecdsa_api.rs +++ b/rust/basic_bitcoin/src/basic_bitcoin/src/ecdsa_api.rs @@ -2,9 +2,22 @@ use ic_cdk::api::management_canister::ecdsa::{ ecdsa_public_key, sign_with_ecdsa, EcdsaCurve, EcdsaKeyId, EcdsaPublicKeyArgument, SignWithEcdsaArgument, }; +use std::cell::RefCell; +use std::collections::HashMap; + +// stores the ecdsa to maintain state across different calls to the canister (not across updates) +thread_local! { + /* flexible */ static ECDSA: RefCell> /*derivation path*/, Vec /*public key*/>>> = RefCell::default(); +} /// Returns the ECDSA public key of this canister at the given derivation path. pub async fn get_ecdsa_public_key(key_name: String, derivation_path: Vec>) -> Vec { + // Retrieve already stored public key + if let Some(key) = ECDSA.with(|ecdsa| { + ecdsa.borrow().as_ref().and_then(|map| map.get(&derivation_path).cloned()) + }) { + return key; + } // Retrieve the public key of this canister at the given derivation path // from the ECDSA API. let canister_id = None; @@ -15,12 +28,23 @@ pub async fn get_ecdsa_public_key(key_name: String, derivation_path: Vec let res = ecdsa_public_key(EcdsaPublicKeyArgument { canister_id, - derivation_path, + derivation_path: derivation_path.clone(), key_id, }) .await; - res.unwrap().0.public_key + let public_key = res.unwrap().0.public_key; + + // Cache the public key + ECDSA.with(|ecdsa| { + let mut map = ecdsa.borrow_mut(); + if map.is_none() { + *map = Some(HashMap::new()); + } + map.as_mut().unwrap().insert(derivation_path, public_key.clone()); + }); + + public_key } pub async fn get_ecdsa_signature(