Add a cache to avoid repeated key derivations

src/common/wallet.c

@@ -424,7 +424,12 @@ int parse_policy_map_key_info(buffer_t *buffer, policy_map_key_info_t *out, int
     return 0;
 }
 
-static int parse_placeholder(buffer_t *in_buf, int version, policy_node_key_placeholder_t *out) {
+// parses a placeholder from in_buf, storing it in out. On success, the pointed placeholder_index is
+// stored in out->placeholder_index, and then it's incremented.
+static int parse_placeholder(buffer_t *in_buf,
+                             int version,
+                             policy_node_key_placeholder_t *out,
+                             uint16_t *placeholder_index) {
     char c;
     if (!buffer_read_u8(in_buf, (uint8_t *) &c) || c != '@') {
         return WITH_ERROR(-1, "Expected key placeholder starting with '@'");
@@ -489,6 +494,8 @@ static int parse_placeholder(buffer_t *in_buf, int version, policy_node_key_plac
         return WITH_ERROR(-1, "Invalid version number");
     }
 
+    out->placeholder_index = *placeholder_index;
+    ++(*placeholder_index);
     return 0;
 }
 
@@ -544,6 +551,15 @@ static int parse_script(buffer_t *in_buf,
                         unsigned int context_flags) {
     int n_wrappers = 0;
 
+    // Keep track of how many key placeholders have been created while parsing
+    // This allows to know the counter even in recursive calls
+    static uint16_t key_placeholder_count = 0;
+
+    if (depth == 0) {
+        // reset the counter on function entry, but not in recursive calls
+        key_placeholder_count = 0;
+    }
+
     policy_node_t *outermost_node = (policy_node_t *) buffer_get_cur(out_buf);
     policy_node_with_script_t *inner_wrapper = NULL;  // pointer to the inner wrapper, if any
 
@@ -1396,7 +1412,7 @@ static int parse_script(buffer_t *in_buf,
 
             node->base.type = token;
 
-            if (0 > parse_placeholder(in_buf, version, key_placeholder)) {
+            if (0 > parse_placeholder(in_buf, version, key_placeholder, &key_placeholder_count)) {
                 return WITH_ERROR(-1, "Couldn't parse key placeholder");
             }
 
@@ -1466,7 +1482,7 @@ static int parse_script(buffer_t *in_buf,
             }
             i_policy_node_key_placeholder(&node->key_placeholder, key_placeholder);
 
-            if (0 > parse_placeholder(in_buf, version, key_placeholder)) {
+            if (0 > parse_placeholder(in_buf, version, key_placeholder, &key_placeholder_count)) {
                 return WITH_ERROR(-1, "Couldn't parse key placeholder");
             }
 
@@ -1606,7 +1622,8 @@ static int parse_script(buffer_t *in_buf,
                     return WITH_ERROR(-1, "Out of memory");
                 }
 
-                if (0 > parse_placeholder(in_buf, version, key_placeholder)) {
+                if (0 >
+                    parse_placeholder(in_buf, version, key_placeholder, &key_placeholder_count)) {
                     return WITH_ERROR(-1, "Error parsing key placeholder");
                 }
 

src/common/wallet.h

@@ -300,6 +300,8 @@ typedef struct {
 
     // common between V1 and V2
     int16_t key_index;  // index of the key
+    int16_t
+        placeholder_index;  // index of the placeholder in the descriptor template, in parsing order
 } policy_node_key_placeholder_t;
 
 DEFINE_REL_PTR(policy_node_key_placeholder, policy_node_key_placeholder_t)

src/handler/lib/policy.c

@@ -469,10 +469,16 @@ __attribute__((warn_unused_result)) static int get_derived_pubkey(
 
     // we derive the /<change>/<address_index> child of this pubkey
     // we reuse the same memory of ext_pubkey
-    bip32_CKDpub(&ext_pubkey,
-                 wdi->change ? key_placeholder->num_second : key_placeholder->num_first,
-                 &ext_pubkey);
-    bip32_CKDpub(&ext_pubkey, wdi->address_index, &ext_pubkey);
+    if (0 > derive_first_step_for_pubkey(&ext_pubkey,
+                                         key_placeholder,
+                                         wdi->sign_psbt_cache,
+                                         wdi->change,
+                                         &ext_pubkey)) {
+        return -1;
+    }
+    if (0 > bip32_CKDpub(&ext_pubkey, wdi->address_index, &ext_pubkey)) {
+        return -1;
+    }
 
     memcpy(out, ext_pubkey.compressed_pubkey, 33);
 

src/handler/lib/policy.h

@@ -2,6 +2,7 @@
 
 #include "../../boilerplate/dispatcher.h"
 #include "../../common/wallet.h"
+#include "../../handler/sign_psbt/sign_psbt_cache.h"
 
 /**
  * Parses a serialized wallet policy, saving the wallet header, the policy map descriptor and the
@@ -48,6 +49,8 @@ typedef struct {
     uint32_t n_keys;        // The number of key information placeholders in the policy
     size_t address_index;   // The address index to use in the derivation
     bool change;            // whether a change address or a receive address is derived
+    sign_psbt_cache_t
+        *sign_psbt_cache;  // If not NULL, the cache for key derivations used during signing
 } wallet_derivation_info_t;
 
 /**