@@ -2,9 +2,22 @@ use ic_cdk::api::management_canister::ecdsa::{
2
2
ecdsa_public_key, sign_with_ecdsa, EcdsaCurve , EcdsaKeyId , EcdsaPublicKeyArgument ,
3
3
SignWithEcdsaArgument ,
4
4
} ;
5
+ use std:: cell:: RefCell ;
6
+ use std:: collections:: HashMap ;
7
+
8
+ // stores the ecdsa to maintain state across different calls to the canister (not across updates)
9
+ thread_local ! {
10
+ /* flexible */ static ECDSA : RefCell <Option <HashMap <Vec <Vec <u8 >>/*derivation path*/ , Vec <u8 > /*public key*/ >>> = RefCell :: default ( ) ;
11
+ }
5
12
6
13
/// Returns the ECDSA public key of this canister at the given derivation path.
7
14
pub async fn get_ecdsa_public_key ( key_name : String , derivation_path : Vec < Vec < u8 > > ) -> Vec < u8 > {
15
+ // Retrieve already stored public key
16
+ if let Some ( key) = ECDSA . with ( |ecdsa| {
17
+ ecdsa. borrow ( ) . as_ref ( ) . and_then ( |map| map. get ( & derivation_path) . cloned ( ) )
18
+ } ) {
19
+ return key;
20
+ }
8
21
// Retrieve the public key of this canister at the given derivation path
9
22
// from the ECDSA API.
10
23
let canister_id = None ;
@@ -15,12 +28,23 @@ pub async fn get_ecdsa_public_key(key_name: String, derivation_path: Vec<Vec<u8>
15
28
16
29
let res = ecdsa_public_key ( EcdsaPublicKeyArgument {
17
30
canister_id,
18
- derivation_path,
31
+ derivation_path : derivation_path . clone ( ) ,
19
32
key_id,
20
33
} )
21
34
. await ;
22
35
23
- res. unwrap ( ) . 0 . public_key
36
+ let public_key = res. unwrap ( ) . 0 . public_key ;
37
+
38
+ // Cache the public key
39
+ ECDSA . with ( |ecdsa| {
40
+ let mut map = ecdsa. borrow_mut ( ) ;
41
+ if map. is_none ( ) {
42
+ * map = Some ( HashMap :: new ( ) ) ;
43
+ }
44
+ map. as_mut ( ) . unwrap ( ) . insert ( derivation_path, public_key. clone ( ) ) ;
45
+ } ) ;
46
+
47
+ public_key
24
48
}
25
49
26
50
pub async fn get_ecdsa_signature (
0 commit comments