Verify that the number of keys in policy being registered is correct

Author: bigspider

bitcoin_client_js/src/__tests__/appClient.test.ts

@@ -267,6 +267,7 @@ describe("test AppClient", () => {
       [
         "[76223a6e/48'/1'/0'/2']tpubDE7NQymr4AFtewpAsWtnreyq9ghkzQBXpCZjWLFVRAvnbf7vya2eMTvT2fPapNqL8SuVvLQdbUbMfWLVDCZKnsEBqp6UK93QEzL8Ck23AwF",
         "[f5acc2fd/48'/1'/0'/2']tpubDFAqEGNyad35aBCKUAXbQGDjdVhNueno5ZZVEn3sQbW5ci457gLR7HyTmHBg93oourBssgUxuWz1jX5uhc1qaqFo9VsybY1J5FuedLfm4dK",
+        "tpubDCoDDpHR1MYXcFrarTcwBufQvWPXSSZpGxjnhRaW612TMxs5TWDEPdbYRHtQdZ9z1UqtKGQKVQ4FqejzbFSdvQvJsD75yrgh7thVoFho6jE",
       ]
     );
 

src/handler/lib/policy.c

@@ -1645,6 +1645,21 @@ int get_key_placeholder_by_index(const policy_node_t *policy,
     return -1;
 }
 
+int count_distinct_keys_info(const policy_node_t *policy) {
+    policy_node_key_placeholder_t placeholder;
+    int ret = -1, cur, n_placeholders;
+
+    for (cur = 0;
+         cur < (n_placeholders = get_key_placeholder_by_index(policy, cur, NULL, &placeholder));
+         ++cur) {
+        if (n_placeholders < 0) {
+            return -1;
+        }
+        ret = MAX(ret, placeholder.key_index + 1);
+    }
+    return ret;
+}
+
 // Utility function to extract and decode the i-th xpub from the keys information vector
 static int get_pubkey_from_merkle_tree(dispatcher_context_t *dispatcher_context,
                                        int wallet_version,

src/handler/lib/policy.h

@@ -195,6 +195,18 @@ __attribute__((warn_unused_result)) int get_key_placeholder_by_index(
     const policy_node_t **out_tapleaf_ptr,
     policy_node_key_placeholder_t *out_placeholder);
 
+/**
+ * Determines the expected number of unique keys in the provided policy's key information.
+ * The function calculates this by finding the maximum key index from placeholders and increments it
+ * by 1. For instance, if the maximum key index found in the placeholders is `n`, then the result
+ * would be `n + 1`.
+ *
+ * @param[in] policy
+ *   Pointer to the root node of the policy
+ * @return the expected number of items in the keys information vector; -1 in case of error.
+ */
+__attribute__((warn_unused_result)) int count_distinct_keys_info(const policy_node_t *policy);
+
 /**
  * Checks if a wallet policy is sane, verifying that pubkeys are never repeated and (if miniscript)
  * that the miniscript is "sane".

src/handler/register_wallet.c

@@ -88,11 +88,18 @@ void handler_register_wallet(dispatcher_context_t *dc, uint8_t protocol_version)
         return;
     }
 
+    if (count_distinct_keys_info(&policy_map.parsed) != (int) wallet_header.n_keys) {
+        PRINTF("Number of keys in descriptor template doesn't provided keys\n");
+        SEND_SW(dc, SW_INCORRECT_DATA);
+        return;
+    }
+
     // Compute the wallet id (sha256 of the serialization)
     get_policy_wallet_id(&wallet_header, wallet_id);
 
     // Verify that the name is acceptable
     if (!is_policy_name_acceptable(wallet_header.name, wallet_header.name_len)) {
+        PRINTF("Policy name is not acceptable\n");
         SEND_SW(dc, SW_INCORRECT_DATA);
         return;
     }

tests/test_register_wallet.py

@@ -117,6 +117,21 @@ def test_register_wallet_invalid_names(client: Client):
             client.register_wallet(wallet)
 
 
+@has_automation("automations/register_wallet_accept.json")
+def test_register_wallet_missing_key(client: Client):
+    wallet = WalletPolicy(
+        name="Missing a key",
+        descriptor_template="wsh(multi(2,@0/**,@1/**))",
+        keys_info=[
+            "[f5acc2fd/48'/1'/0'/2']tpubDFAqEGNyad35aBCKUAXbQGDjdVhNueno5ZZVEn3sQbW5ci457gLR7HyTmHBg93oourBssgUxuWz1jX5uhc1qaqFo9VsybY1J5FuedLfm4dK",
+            # the second key is missing
+        ],
+    )
+
+    with pytest.raises(IncorrectDataError):
+        client.register_wallet(wallet)
+
+
 @has_automation("automations/register_wallet_accept.json")
 def test_register_wallet_unsupported_policy(client: Client):
     # valid policies, but not supported (might change in the future)