From c586ea0a81a50ba049022d487a38c06755955684 Mon Sep 17 00:00:00 2001
From: Daniils Tolmacovs <tolmadan@Daniilss-MacBook-Air.local>
Date: Wed, 16 Oct 2024 16:47:45 +0300
Subject: [PATCH] feat(basic_bitcoin): cache ecdsa

---
 .../src/basic_bitcoin/src/ecdsa_api.rs        | 28 +++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

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<Option<HashMap<Vec<Vec<u8>> /*derivation path*/, Vec<u8> /*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<u8>>) -> Vec<u8> {
+    // 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<Vec<u8>
 
     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(