descriptor: allow iterating keys

Author: jgriffiths

include/wally.hpp

@@ -504,6 +504,12 @@ inline int descriptor_get_features(const DESCRIPTOR& descriptor, uint32_t* value
     return ret;
 }
 
+template <class DESCRIPTOR>
+inline int descriptor_get_key(const DESCRIPTOR& descriptor, size_t index, char** output) {
+    int ret = ::wally_descriptor_get_key(detail::get_p(descriptor), index, output);
+    return ret;
+}
+
 template <class DESCRIPTOR>
 inline int descriptor_get_network(const DESCRIPTOR& descriptor, uint32_t* value_out) {
     int ret = ::wally_descriptor_get_network(detail::get_p(descriptor), value_out);

include/wally_descriptor.h

@@ -131,18 +131,6 @@ WALLY_CORE_API int wally_descriptor_get_features(
     const struct wally_descriptor *descriptor,
     uint32_t *value_out);
 
-/**
- * Get the number of keys in a parsed output descriptor or miniscript expression.
- *
- * :param descriptor: Parsed output descriptor or miniscript expression.
- * :param value_out: Destination for the number of keys.
- *
-* .. note:: Repeated keys are counted once for each time they appear.
- */
-WALLY_CORE_API int wally_descriptor_get_num_keys(
-    const struct wally_descriptor *descriptor,
-    uint32_t *value_out);
-
 /**
  * Get the number of variants in a parsed output descriptor or miniscript expression.
  *
@@ -193,6 +181,31 @@ WALLY_CORE_API int wally_descriptor_get_depth(
     const struct wally_descriptor *descriptor,
     uint32_t *value_out);
 
+/**
+ * Get the number of keys in a parsed output descriptor or miniscript expression.
+ *
+ * :param descriptor: Parsed output descriptor or miniscript expression.
+ * :param value_out: Destination for the number of keys.
+ *
+* .. note:: Repeated keys are counted once for each time they appear.
+ */
+WALLY_CORE_API int wally_descriptor_get_num_keys(
+    const struct wally_descriptor *descriptor,
+    uint32_t *value_out);
+
+/**
+ * Get the string representation of a key in a parsed output descriptor or miniscript expression.
+ *
+ * :param descriptor: Parsed output descriptor or miniscript expression.
+ * :param index: The zero-based index of the key to get.
+ * :param output: Destination for the resulting string representation.
+ *|    The string returned should be freed using `wally_free_string`.
+ */
+WALLY_CORE_API int wally_descriptor_get_key(
+    const struct wally_descriptor *descriptor,
+    size_t index,
+    char **output);
+
 /**
  * Get the maximum length of a script corresponding to an output descriptor.
  *

src/descriptor.c

@@ -206,7 +206,7 @@ static int ctx_add_key_node(ms_ctx *ctx, ms_node *node)
 {
     const char *v = (char *)node;
     return map_add(&ctx->keys, NULL, ctx->keys.num_items,
-                   (unsigned char *)v, 1, false, false);
+                   (unsigned char *)v, 1, true, false);
 }
 
 /* Built-in miniscript expressions */
@@ -2840,17 +2840,6 @@ int wally_descriptor_get_features(const struct wally_descriptor *descriptor,
                              offsetof(struct wally_descriptor, features));
 }
 
-int wally_descriptor_get_num_keys(const struct wally_descriptor *descriptor,
-                                  uint32_t *value_out)
-{
-    if (value_out)
-        *value_out = 0;
-    if (!descriptor || !value_out)
-        return WALLY_EINVAL;
-    *value_out = (uint32_t)descriptor->keys.num_items;
-    return WALLY_OK;
-}
-
 int wally_descriptor_get_num_variants(const struct wally_descriptor *descriptor,
                                       uint32_t *value_out)
 {
@@ -2887,3 +2876,50 @@ int wally_descriptor_get_depth(const struct wally_descriptor *descriptor,
     *value_out = node_get_depth(descriptor->top_node) - 1;
     return WALLY_OK;
 }
+
+int wally_descriptor_get_num_keys(const struct wally_descriptor *descriptor,
+                                  uint32_t *value_out)
+{
+    if (value_out)
+        *value_out = 0;
+    if (!descriptor || !value_out)
+        return WALLY_EINVAL;
+    *value_out = (uint32_t)descriptor->keys.num_items;
+    return WALLY_OK;
+}
+
+/* Ignore incorrect warnings from the ms_node cast below */
+#pragma GCC diagnostic ignored "-Wcast-align"
+#if defined(__clang__)
+#pragma clang diagnostic ignored "-Wcast-align"
+#endif
+
+int wally_descriptor_get_key(const struct wally_descriptor *descriptor,
+                             size_t index, char **output)
+{
+    const ms_node *node;
+
+    if (output)
+        *output = 0;
+    if (!descriptor || index >= descriptor->keys.num_items || !output)
+        return WALLY_EINVAL;
+
+    node = (ms_node *)descriptor->keys.items[index].value;
+    if (node->kind == KIND_PUBLIC_KEY) {
+        return wally_hex_from_bytes((const unsigned char *)node->data,
+                                    node->data_len, output);
+    }
+    if (node->kind == KIND_PRIVATE_KEY) {
+        uint32_t flags = node->flags & NF_IS_UNCOMPRESSED ? WALLY_WIF_FLAG_UNCOMPRESSED : 0;
+        if (!descriptor->addr_ver)
+            return WALLY_EINVAL; /* Must have a network to fetch private keys */
+        return wally_wif_from_bytes((const unsigned char *)node->data, node->data_len,
+                                    descriptor->addr_ver->version_wif,
+                                    flags, output);
+    }
+    if ((node->kind & KIND_BIP32) != KIND_BIP32)
+        return WALLY_ERROR; /* Unknown key type, should not happen */
+    if (!(*output = wally_strdup_n(node->data, node->data_len)))
+        return WALLY_ENOMEM;
+    return WALLY_OK;
+}

src/swig_java/swig.i

@@ -561,6 +561,7 @@ static jobjectArray create_jstringArray(JNIEnv *jenv, char **p, size_t len) {
 %returns_size_t(wally_descriptor_get_depth);
 %returns_size_t(wally_descriptor_get_features);
 %returns_size_t(wally_descriptor_get_network);
+%returns_string(wally_descriptor_get_key);
 %returns_size_t(wally_descriptor_get_num_keys);
 %returns_size_t(wally_descriptor_get_num_paths);
 %returns_size_t(wally_descriptor_get_num_variants);

src/test/test_descriptor.py

@@ -314,5 +314,35 @@ def make_keys(xpubs):
             self.assertEqual(ret, WALLY_EINVAL)
             wally_map_free(keys)
 
+    def test_key_iteration(self):
+        """Test iterating descriptor keys"""
+        k1 = 'xpub661MyMwAqRbcFW31YEwpkMuc5THy2PSt5bDMsktWQcFF8syAmRUapSCGu8ED9W6oDMSgv6Zz8idoc4a6mr8BDzTJY47LJhkJ8UB7WEGuduB'
+        k2 = 'xprvA2YKGLieCs6cWCiczALiH1jzk3VCCS5M1pGQfWPkamCdR9UpBgE2Gb8AKAyVjKHkz8v37avcfRjdcnP19dVAmZrvZQfvTcXXSAiFNQ6tTtU'
+        wif = 'L1AAHuEC7XuDM7pJ7yHLEqYK1QspMo8n1kgxyZVdgvEpVC1rkUrM'
+        pk = '03b428da420cd337c7208ed42c5331ebb407bb59ffbe3dc27936a227c619804284'
+        pk_u = '0414fc03b8df87cd7b872996810db8458d61da8448e531569c8517b469a119d267be5645686309c6e6736dbd93940707cc9143d3cf29f1b877ff340e2cb2d259cf'
+        # Valid args
+        for descriptor, expected in [
+            # Bip32 xpub
+            (f'pkh({k1})', k1),
+            # BIP32 xprv
+            (f'pkh({k2}/*)', k2),
+            # WIF
+            (f'pkh({wif})', wif),
+            # Hex pubkey, compressed
+            (f'pk({pk})', pk),
+            # Hex pubkey, uncompressed
+            (f'pk({pk_u})', pk_u),
+        ]:
+            d = c_void_p()
+            ret = wally_descriptor_parse(descriptor, None, NETWORK_BTC_MAIN, 0, d)
+            self.assertEqual(ret, WALLY_OK)
+            ret, num_keys = wally_descriptor_get_num_keys(d)
+            self.assertEqual((ret, num_keys), (WALLY_OK, 1))
+            ret, key_str = wally_descriptor_get_key(d, 0)
+            self.assertEqual((ret, key_str), (WALLY_OK, expected))
+            wally_descriptor_free(d)
+
+
 if __name__ == '__main__':
     unittest.main()

src/test/util.py

@@ -313,6 +313,7 @@ class wally_psbt(Structure):
     ('wally_descriptor_get_checksum', c_int, [c_void_p, c_uint32, c_char_p_p]),
     ('wally_descriptor_get_depth', c_int, [c_void_p, c_uint32_p]),
     ('wally_descriptor_get_features', c_int, [c_void_p, c_uint32_p]),
+    ('wally_descriptor_get_key', c_int, [c_void_p, c_size_t, c_char_p_p]),
     ('wally_descriptor_get_network', c_int, [c_void_p, c_uint32_p]),
     ('wally_descriptor_get_num_keys', c_int, [c_void_p, c_uint32_p]),
     ('wally_descriptor_get_num_paths', c_int, [c_void_p, c_uint32_p]),

src/wasm_package/src/functions.js

@@ -123,6 +123,7 @@ export function descriptor_free(descriptor: Ref_wally_descriptor): void;
 export function descriptor_get_checksum(descriptor: Ref_wally_descriptor, flags: number): string;
 export function descriptor_get_depth(descriptor: Ref_wally_descriptor): number;
 export function descriptor_get_features(descriptor: Ref_wally_descriptor): number;
+export function descriptor_get_key(descriptor: Ref_wally_descriptor, index: number): string;
 export function descriptor_get_network(descriptor: Ref_wally_descriptor): number;
 export function descriptor_get_num_keys(descriptor: Ref_wally_descriptor): number;
 export function descriptor_get_num_paths(descriptor: Ref_wally_descriptor): number;

src/wasm_package/src/index.d.ts

@@ -81,6 +81,7 @@ EXPORTED_FUNCTIONS="['_malloc','_free','_bip32_key_free' \
 ,'_wally_descriptor_get_checksum' \
 ,'_wally_descriptor_get_depth' \
 ,'_wally_descriptor_get_features' \
+,'_wally_descriptor_get_key' \
 ,'_wally_descriptor_get_network' \
 ,'_wally_descriptor_get_num_keys' \
 ,'_wally_descriptor_get_num_paths' \